def METHOD_NAME(self, device: Optional[str] = None) -> nn.Module:
    """Since teacher_model is registered as a plain object, it is necessary
    to put the teacher model to other device when calling ``to``
    function."""
    self.teacher_model.METHOD_NAME(device=device)
    return super().METHOD_NAME(device=device)