Unverified 提交 b2194b90 authored 作者: Glenn Jocher's avatar Glenn Jocher 提交者: GitHub

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

上级 7830e91b
......@@ -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
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论