update save statistic information for era5_dataset.py

video_prediction_tools/model_modules/video_prediction/datasets/era5_dataset.py

@@ -21,6 +21,7 @@ class ERA5Dataset(BaseDataset):
                          shuffled_on_val = shuffled_on_val, output_dir = output_dir)
 
         self.shuffled_on_val = shuffled_on_val
+        self.data_arr, self.init_times = self.load_data(self.filenames)
 
     def specific_hparams(self)-> list:
         s_hparams = ["context_frames", "sequence_length", "shift"]
@@ -82,7 +83,6 @@ class ERA5Dataset(BaseDataset):
         """
 
         shuffle = self.mode == 'train' or (self.mode == 'val' and self.shuffle_on_val)
-        data_arr, init_times = self.load_data(self.filenames)
         
         def normalize_fn(x:tf.Tensor, min_value:float, max_value:float):
             return tf.divide(tf.subtract(x, min_value), tf.subtract(max_value, min_value))
@@ -121,7 +121,7 @@ class ERA5Dataset(BaseDataset):
             raise ("The filenames list is empty for {} dataset, please make sure your data_split dictionary is configured correctly".format(self.mode))
         else:
             #group the data into sequenceds
-            dataset = tf.data.Dataset.from_generator(lambda: iter(zip(data_arr, init_times)),
+            dataset = tf.data.Dataset.from_generator(lambda: iter(zip(self.data_arr, self.init_times)),
                                                      output_types=(tf.float32,tf.string), output_shapes=((self.nlat,self.nlon,self.n_vars),()))
             dataset = dataset.window(self.sequence_length, shift = self.shift, drop_remainder=True)
             dataset = dataset.flat_map(lambda x,y: tf.data.Dataset.zip((x.batch(self.sequence_length),y.batch(self.sequence_length))))