`yolo.py --profile` default GPU batch size 16
Browse files- models/yolo.py +1 -1
models/yolo.py
CHANGED
@@ -314,7 +314,7 @@ if __name__ == '__main__':
|
|
314 |
|
315 |
# Profile
|
316 |
if opt.profile:
|
317 |
-
img = torch.rand(
|
318 |
y = model(img, profile=True)
|
319 |
|
320 |
# Test all models
|
|
|
314 |
|
315 |
# Profile
|
316 |
if opt.profile:
|
317 |
+
img = torch.rand(16 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
|
318 |
y = model(img, profile=True)
|
319 |
|
320 |
# Test all models
|