Try caching dataset.

video_prediction_tools/model_modules/video_prediction/datasets/era5_dataset.py

@@ -181,6 +181,8 @@ class ERA5Dataset(object):
             random.shuffle(filenames)
         dataset = tf.data.TFRecordDataset(filenames, buffer_size = 8* 1024 * 1024) 
        
+        # cache dataset
+        dataset = dataset.cache()  
         # dataset = dataset.filter(self.filter)
         if shuffle:
             dataset = dataset.apply(tf.contrib.data.shuffle_and_repeat(buffer_size =1024, count = self.num_epochs))