Update TorchScript suffix to *.torchscript (#5856)
This commit is contained in:
1 parent
30db14fea8
commit
00e308f7be
5 files changed
+20
-20
No files matched your search
@@ -81,18 +81,18 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
|
||||
imgsz = check_img_size(imgsz, s=stride) # check image size
|
||||
|
||||
# Half
|
||||
half &= (pt or engine) and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
|
||||
if pt:
|
||||
half &= (pt or jit or engine) and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
|
||||
if pt or jit:
|
||||
model.model.half() if half else model.model.float()
|
||||
|
||||
# Dataloader
|
||||
if webcam:
|
||||
view_img = check_imshow()
|
||||
cudnn.benchmark = True # set True to speed up constant image size inference
|
||||
dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt and not jit)
|
||||
dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt)
|
||||
bs = len(dataset) # batch_size
|
||||
else:
|
||||
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt and not jit)
|
||||
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt)
|
||||
bs = 1 # batch_size
|
||||
vid_path, vid_writer = [None] * bs, [None] * bs
|
||||
|
||||
|
||||
Reference in new issue
Block a user