From 5fb661b7d43021f841bf84d6c10eed69970fa257 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Mon, 9 Mar 2020 13:33:23 -0700
Subject: [PATCH] updates

---
 test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test.py b/test.py
index 2b6b8f23..f6a68859 100644
--- a/test.py
+++ b/test.py
@@ -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