Updated a bug, added batches_per_epoch method to the dataset

credit/datasets/era5_multistep_batcher.py

@@ -1,5 +1,6 @@
 import logging
 import time
+import math
 import queue
 import multiprocessing
 from threading import Thread
@@ -332,8 +333,12 @@ def set_epoch(self, epoch):
         self.current_epoch = epoch
         self.sampler.set_epoch(epoch)
         self.batch_indices = list(self.sampler)
+        self.batch_call_count = 0
         self.initialize_batch()
 
+    def batches_per_epoch(self):
+        return math.ceil((self.batch_indices) / self.batch_size)
+
     def __getitem__(self, _):
         """
         Fetches the current forecast step data for each item in the batch.
@@ -507,6 +512,7 @@ def set_epoch(self, epoch):
         self.current_epoch = epoch
         self.sampler.set_epoch(epoch)
         self.batch_indices = list(self.sampler)
+        self.batch_call_count = 0
         self.initialize_batch()
 
         # Start prefetching thread