BucketedBatchSampler的batch_id_in_epoch实现

fastNLP/core/samplers/reproducible_batch_sampler.py

@@ -411,4 +411,12 @@ class BucketedBatchSampler(ReproducibleBatchSampler):
         self.old_num_replicas = states['num_replicas']
 
     def set_epoch(self, epoch):
-        self.epoch = epoch
+        self.epoch = epoch
+
+    @property
+    def batch_idx_in_epoch(self):
+        if self.drop_last:
+            return len(self.dataset) // self.batch_size - (len(self.dataset) - self.num_consumed_samples) // self.batch_size
+        else:
+            return (len(self.dataset) + self.batch_size - 1) // self.batch_size - \
+                   (len(self.dataset) - self.num_consumed_samples + self.batch_size - 1) // self.batch_size