diff --git a/modules/deeplearning/cloudheight.py b/modules/deeplearning/cloudheight.py
index 55585f0e53a9f84e5d87dec611c95e5575f3b162..c0fe2d60737b0e79903566da84933ff771333a0c 100644
--- a/modules/deeplearning/cloudheight.py
+++ b/modules/deeplearning/cloudheight.py
@@ -14,7 +14,6 @@ LOG_DEVICE_PLACEMENT = False
 
 CACHE_DATA_IN_MEM = True
 CACHE_GFS = True
-DISK_CACHE = True
 
 PROC_BATCH_SIZE = 40
 PROC_BATCH_BUFFER_SIZE = 50000
@@ -242,6 +241,8 @@ class CloudHeightNN:
         self.inputs.append(self.X_prof)
         self.inputs.append(self.X_sfc)
 
+        self.DISK_CACHE = True
+
         if datapath is not None:
             self.DISK_CACHE = False
             f = open(datapath, 'rb')
@@ -800,7 +801,7 @@ class CloudHeightNN:
             print('loss, acc: ', self.test_loss.result(), self.test_accuracy.result())
             ckpt_manager.save()
 
-            if DISK_CACHE and epoch == 0:
+            if self.DISK_CACHE and epoch == 0:
                 f = open(cachepath, 'wb')
                 pickle.dump(self.in_mem_data_cache, f)
                 f.close()