Fixed CPU context on segmentation train (#485)

Fix CPU context
dmlc · Nov 27, 2018 · 87cddd7 · 87cddd7
1 parent 48ea65f
commit 87cddd7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/scripts/segmentation/train.py b/scripts/segmentation/train.py
@@ -91,7 +91,7 @@ def parse_args():
     if args.no_cuda:
         print('Using CPU')
         args.kvstore = 'local'
-        args.ctx = mx.cpu(0)
+        args.ctx = [mx.cpu(0)]
     else:
         print('Number of GPUs:', args.ngpus)
         args.ctx = [mx.gpu(i) for i in range(args.ngpus)]