glenn-jocher commited on
Commit
67d5e65
1 Parent(s): 2703ac7

profiling update

Browse files
Files changed (1) hide show
  1. models/yolo.py +2 -2
models/yolo.py CHANGED
@@ -180,8 +180,8 @@ if __name__ == '__main__':
180
  model.train()
181
 
182
  # Profile
183
- img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
184
- y = model(img, profile=True)
185
  # print([y[0].shape] + [x.shape for x in y[1]])
186
 
187
  # ONNX export
 
180
  model.train()
181
 
182
  # Profile
183
+ # img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
184
+ # y = model(img, profile=True)
185
  # print([y[0].shape] + [x.shape for x in y[1]])
186
 
187
  # ONNX export