Ver código fonte

Multi-GPU disabled in test.py

5.0
Glenn Jocher 4 anos atrás
pai
commit
72d06147b2
1 arquivos alterados com 3 adições e 3 exclusões
  1. +3
    -3
      test.py

+ 3
- 3
test.py Ver arquivo

@@ -1,7 +1,6 @@
import argparse
import json

import yaml
from torch.utils.data import DataLoader

from utils.datasets import *
@@ -40,8 +39,9 @@ def test(data,
if half:
model.half() # to FP16

if device.type != 'cpu' and torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
# Multi-GPU disabled, incompatible with .half()
# if device.type != 'cpu' and torch.cuda.device_count() > 1:
# model = nn.DataParallel(model)

else: # called by train.py
training = True

Carregando…
Cancelar
Salvar