make some changes for running on SLURM

modules/deeplearning/icing_cnn.py

@@ -219,6 +219,7 @@ class IcingIntensityNN:
         tf.debugging.set_log_device_placement(LOG_DEVICE_PLACEMENT)
 
         # Note: Don't do this anymore, because nobody else willing to do so as well!
+        # Also, doesn't seem to play well with SLURM
         # gpus = tf.config.experimental.list_physical_devices('GPU')
         # if gpus:
         #     try:
@@ -918,12 +919,13 @@ class IcingIntensityNN:
         self.test_preds = preds
 
     def run(self, filename_l1b_trn, filename_l1b_tst, filename_l2_trn, filename_l2_tst):
-        with tf.device('/device:GPU:'+str(self.gpu_device)):
-            self.setup_pipeline(filename_l1b_trn, filename_l1b_tst, filename_l2_trn, filename_l2_tst)
-            self.build_model()
-            self.build_training()
-            self.build_evaluation()
-            self.do_training()
+        # This doesn't really play well with SLURM
+        # with tf.device('/device:GPU:'+str(self.gpu_device)):
+        self.setup_pipeline(filename_l1b_trn, filename_l1b_tst, filename_l2_trn, filename_l2_tst)
+        self.build_model()
+        self.build_training()
+        self.build_evaluation()
+        self.do_training()
 
     def run_restore(self, filename_l1b, filename_l2, ckpt_dir):
         self.setup_test_pipeline(filename_l1b, filename_l2)