def METHOD_NAME(self, model: nn.Module, optimizer: OptimizerWrapper) -> Iterator[None]:
    assert isinstance(model, TorchDDPModel), 'Model is not boosted by TorchDDPPlugin.'
    return model.module.METHOD_NAME()