Method train

colossalai/legacy/engine/_base_engine.py:210–213

Sets the model to training mode.

(self)

        return output, label, loss

    def train(self):
        """Sets the model to training mode."""
        self.training = True
        self._model.train()

    def eval(self):
        """Sets the model to evaluation mode."""