diff --git a/modules/deeplearning/icing_cnn.py b/modules/deeplearning/icing_cnn.py index 7977e1adde7a597888cf81ac9d7abf62c0d5a2e0..87a380c3509ca3f2e063489b5b63b622e39bf1ab 100644 --- a/modules/deeplearning/icing_cnn.py +++ b/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)