Browse Source

`yolo.py --profile` default GPU batch size 16

modifyDataloader
Glenn Jocher GitHub 2 years ago
parent
commit
7830e91b9a
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
  1. +1
    -1
      models/yolo.py

+ 1
- 1
models/yolo.py View File

@@ -314,7 +314,7 @@ if __name__ == '__main__':

# Profile
if opt.profile:
img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
img = torch.rand(16 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
y = model(img, profile=True)

# Test all models

Loading…
Cancel
Save