updates

2020-03-09 13:33:23 -07:00
parent 6bd51b75ea
commit 5fb661b7d4
1 changed files with 1 additions and 1 deletions
@@ -38,7 +38,7 @@ def test(cfg,
        else:  # darknet format
            load_darknet_weights(model, weights)

-        if torch.cuda.device_count() > 1:
+        if device.type != 'cpu' and torch.cuda.device_count() > 1:
            model = nn.DataParallel(model)
    else:  # called by train.py
        device = next(model.parameters()).device  # get model device