New smart_inference_mode() conditional decorator (#8957)

New smart_inference_mode()
This commit is contained in:
Glenn Jocher authored and GitHub committed 2022-08-13 20:38:51 +02:00
1 parent 6aed0a7c00
commit dc38cd03f4
6 files changed
+26 -18

No files matched your search

+2 -2
View File
@@ -69,7 +69,7 @@ from models.yolo import Detect
from utils.dataloaders import LoadImages
from utils.general import (LOGGER, check_dataset, check_img_size, check_requirements, check_version, check_yaml,
colorstr, file_size, print_args, url2file)
from utils.torch_utils import select_device
from utils.torch_utils import select_device, smart_inference_mode
def export_formats():
@@ -455,7 +455,7 @@ def export_tfjs(file, prefix=colorstr('TensorFlow.js:')):
LOGGER.info(f'\n{prefix} export failure: {e}')
@torch.no_grad()
@smart_inference_mode()
def run(
data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
weights=ROOT / 'yolov5s.pt', # weights path