Parcourir la source

update test.py model.half()

5.0
Glenn Jocher il y a 4 ans
Parent
révision
0afbb8d498
1 fichiers modifiés avec 2 ajouts et 2 suppressions
  1. +2
    -2
      test.py

+ 2
- 2
test.py Voir le fichier

@@ -41,9 +41,9 @@ def test(data,
# model = nn.DataParallel(model)

# Half
half = device.type != 'cpu' and torch.cuda.device_count() == 1 # half precision only supported on single-GPU
half = device.type != 'cpu' # half precision only supported on CUDA
if half:
model.half() # to FP16
model.half()

# Configure
model.eval()

Chargement…
Annuler
Enregistrer