diff --git a/train.py b/train.py
index 57b5d44b..099730f3 100644
--- a/train.py
+++ b/train.py
@@ -48,8 +48,8 @@ def train(
         # Load weights to resume from
         model.load_state_dict(checkpoint['model'])
 
-        if torch.cuda.device_count() > 1:
-            model = nn.DataParallel(model)
+        # if torch.cuda.device_count() > 1:
+        #   model = nn.DataParallel(model)
         model.to(device).train()
 
         # Transfer learning (train only YOLO layers)
@@ -75,8 +75,8 @@ def train(
             load_darknet_weights(model, weights + 'yolov3-tiny.conv.15')
             cutoff = 15
 
-        if torch.cuda.device_count() > 1:
-            model = nn.DataParallel(model)
+        # if torch.cuda.device_count() > 1:
+        #    model = nn.DataParallel(model)
         model.to(device).train()
 
         # Set optimizer