From 209a1c4cb2a0f1bf737f623295ed1ee5da846adc Mon Sep 17 00:00:00 2001
From: Ahmed Ibrahim <nady@vt.edu>
Date: Fri, 6 Jul 2018 12:20:38 -0400
Subject: [PATCH] fix a bug in CPU mode

DataParallel should be used only with GPUs.
If a user disabled Cuda in order to run the model on CPU, an error will show up.
---
 train.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/train.py b/train.py
index a5eb890..eaa804d 100644
--- a/train.py
+++ b/train.py
@@ -141,8 +141,10 @@ def main():
     print("build vnet")
     model = vnet.VNet(elu=False, nll=nll)
     batch_size = args.ngpu*args.batchSz
-    gpu_ids = range(args.ngpu)
-    model = nn.parallel.DataParallel(model, device_ids=gpu_ids)
+    # Use DataParallel only if Cuda is enabled.
+    if args.cuda:
+        gpu_ids = range(args.ngpu)
+        model = nn.parallel.DataParallel(model, device_ids=gpu_ids)
 
     if args.resume:
         if os.path.isfile(args.resume):