glenn-jocher commited on
Commit
7830e91
·
unverified ·
1 Parent(s): 3373aab

`yolo.py --profile` default GPU batch size 16

Browse files
Files changed (1) hide show
  1. 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(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
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