Multi-GPU disabled in test.py

This commit is contained in:
Glenn Jocher committed 2020-06-16 14:42:30 -07:00
1 parent bafbc65ee3
commit 72d06147b2
1 file changed
+3 -3
+3 -3
View File
@@ -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