1.修复TrainBatchLoop和EvaluateBatchLoop中的问题; 2.修复OverfitDataLoader不break的问题

Author: yhcc

fastNLP/core/controllers/loops/evaluate_batch_loop.py

@@ -18,6 +18,8 @@ class EvaluateBatchLoop(Loop):
     def __init__(self, batch_step_fn:Optional[Callable]=None):
         if batch_step_fn is not None:
             self.batch_step_fn = batch_step_fn
+        else:
+            self.batch_step_fn = EvaluateBatchLoop.batch_step_fn
 
     def run(self, evaluator, dataloader) -> Dict:
         r"""

fastNLP/core/controllers/loops/train_batch_loop.py

@@ -20,6 +20,8 @@ class TrainBatchLoop(Loop):
     def __init__(self, batch_step_fn: Optional[Callable] = None):
         if batch_step_fn is not None:
             self.batch_step_fn = batch_step_fn
+        else:
+            self.batch_step_fn = TrainBatchLoop.batch_step_fn
 
     def run(self, trainer, dataloader):
         r"""

fastNLP/core/dataloaders/utils.py

@@ -131,6 +131,8 @@ def __init__(self, dataloader, overfit_batches: int, batches=None):
             for idx, batch in enumerate(dataloader):
                 if idx < self.overfit_batches or self.overfit_batches <= -1:
                     self.batches.append(batch)
+                else:
+                    break
         else:
             assert isinstance(batches, list)
             self.batches = batches