瀏覽代碼

Multi-GPU disabled in test.py

5.0
Glenn Jocher 4 年之前
父節點
當前提交
72d06147b2
共有 1 個檔案被更改,包括 3 行新增3 行删除
  1. +3
    -3
      test.py

+ 3
- 3
test.py 查看文件

@@ -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

Loading…
取消
儲存