small fix cpu only version

thomwolf · thomwolf · commit 37e77aff1943 · 2018-06-28T14:47:01.000+02:00
diff --git a/train.py b/train.py
@@ -260,7 +260,7 @@ def run_epoch():
 
     n_train = len(trY)
     n_valid = len(vaY)
-    n_batch_train = args.n_batch * n_gpu
+    n_batch_train = args.n_batch * max(n_gpu, 1)
     n_updates_total = (n_train // n_batch_train) * args.n_iter
 
     if n_gpu > 1:
