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

`yolo.py --profile` updates (#7170)

Browse files
Files changed (1) hide show
  1. models/yolo.py +2 -2
models/yolo.py CHANGED
@@ -310,11 +310,11 @@ if __name__ == '__main__':
310
 
311
  # Create model
312
  model = Model(opt.cfg).to(device)
313
- model.train()
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
 
310
 
311
  # Create model
312
  model = Model(opt.cfg).to(device)
 
313
 
314
  # Profile
315
  if opt.profile:
316
+ model.eval().fuse()
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