diff --git a/modules/deeplearning/cloudheight.py b/modules/deeplearning/cloudheight.py
index b98c781c91252fc7f599fc9a1d305144341d31a8..5e9f7fe716cc91391429e415fcf306be501322c5 100644
--- a/modules/deeplearning/cloudheight.py
+++ b/modules/deeplearning/cloudheight.py
@@ -845,9 +845,10 @@ class CloudHeightNN:
 
     def run(self, matchup_dict, train_dict=None, valid_dict=None):
         #with self.strategy.scope():
-        self.setup_pipeline(matchup_dict, train_dict=train_dict, valid_test_dict=valid_dict)
-        self.build_model()
-        self.build_training()
+        with tf.device('/physical_device:GPU:2'):
+            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()