Make sure train doesn't crash when called at max_epoch (#608)

2026-09-09 11:32:07 +08:00 · 2019-12-07 23:22:03 -05:00
parent 4970624f8b
commit 2baa80d626
1 changed files with 1 additions and 0 deletions
@@ -258,6 +258,7 @@ class TrainerTrainLoopMixin(ABC):
        pass

    def train(self):
+        model = self.get_model()
        # run all epochs
        for epoch in range(self.current_epoch, self.max_epochs):
            # set seed for distributed sampler (enables shuffling for each epoch)