diff --git a/modules/deeplearning/cloudheight.py b/modules/deeplearning/cloudheight.py
index 5e9f7fe716cc91391429e415fcf306be501322c5..44b7088c207e8eea6fb04e2370614425ea8246b7 100644
--- a/modules/deeplearning/cloudheight.py
+++ b/modules/deeplearning/cloudheight.py
@@ -250,7 +250,6 @@ class CloudHeightNN:
         tf.debugging.set_log_device_placement(LOG_DEVICE_PLACEMENT)
 
         gpus = tf.config.experimental.list_physical_devices('GPU')
-        print(gpus)
         if gpus:
             try:
                 # Currently, memory growth needs to be the same across GPUs
@@ -845,12 +844,12 @@ class CloudHeightNN:
 
     def run(self, matchup_dict, train_dict=None, valid_dict=None):
         #with self.strategy.scope():
-        with tf.device('/physical_device:GPU:2'):
+        with tf.device('/device:GPU:'+self.gpu_device):
             self.setup_pipeline(matchup_dict, train_dict=train_dict, valid_test_dict=valid_dict)
             self.build_model()
             self.build_training()
-        self.build_evaluation()
-        self.do_training()
+            self.build_evaluation()
+            self.do_training()
 
     def run_restore(self, matchup_dict, ckpt_dir):
         self.setup_pipeline(None, None, matchup_dict)