Glenn Jocher
GitHub
2 years ago
1 changed files with 2 additions and 2 deletions
+ 2
- 2
models/yolo.py
View File
| # Create model | # Create model | ||||
| model = Model(opt.cfg).to(device) | model = Model(opt.cfg).to(device) | ||||
| model.train() | |||||
| # Profile | # Profile | ||||
| if opt.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) | y = model(img, profile=True) | ||||
| # Test all models | # Test all models |
Loading…