diff --git a/training_cli.py b/training_cli.py
index a2b30a5..39e042f 100644
--- a/training_cli.py
+++ b/training_cli.py
@@ -6,9 +6,12 @@
 
 from lib.train import utils
 import datetime
-
+#  Fix for training won't start with 2 GPUs, even if only 1 enabled
 hps = utils.get_hparams()
 os.environ["CUDA_VISIBLE_DEVICES"] = hps.gpus.replace("-", ",")
+n_gpus = len(hps.gpus.split("-"))
+from random import shuffle, randint
+import traceback, json, argparse, itertools, math, torch, pdb
 os.environ["NCCL_P2P_DISABLE"] = 1
 from random import shuffle, randint