`yolo.py --profile` updates (#7170)
Browse files- 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 |
-
|
|
|
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
|