def METHOD_NAME(self,
              model: nn.Module,
              optimizer: Optional[Optimizer] = None,
              criterion: Optional[Callable] = None,
              ) -> Tuple[nn.Module, OptimizerWrapper, Callable]:
    model = TorchAMPModule(model)
    if optimizer is not None:
        optimizer = TorchAMPOptimizer(optimizer, **self.torch_amp_kwargs)
    if criterion is not None:
        criterion = TorchAMPModule(criterion)
    return model, optimizer, criterion