yolo.py --profile updates (#7170)

This commit is contained in:
Glenn Jocher authored and GitHub committed 2022-03-27 20:24:42 +02:00
1 parent 7830e91b9a
commit b2194b9015
1 file changed
+2 -2
+2 -2
View File
@@ -310,11 +310,11 @@ if __name__ == '__main__':
# Create model
model = Model(opt.cfg).to(device)
model.train()
# Profile
if opt.profile:
img = torch.rand(16 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
model.eval().fuse()
img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
y = model(img, profile=True)
# Test all models