snapshot...

modules/deeplearning/cloudheight.py

@@ -260,6 +260,8 @@ class CloudHeightNN:
                 # Memory growth must be set before GPUs have been initialized
                 print(e)
 
+        self.strategy = tf.distribute.MirroredStrategy()
+
     def get_in_mem_data_batch(self, time_keys):
         images = []
         vprof = []
@@ -808,6 +810,7 @@ class CloudHeightNN:
         self.writer_valid.close()
 
     def build_model(self):
+        with self.strategy.scope():
             flat = self.build_cnn()
             flat_1d = self.build_1d_cnn()
             # flat_anc = self.build_anc_dnn()