yolo.py --profile updates (#7170)
This commit is contained in:
1 parent
7830e91b9a
commit
b2194b9015
1 file changed
+2
-2
+2
-2
@@ -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
|
||||
|
||||
Reference in new issue
Block a user