def METHOD_NAME(self, module: torch.nn.Module, input, output):
    self.adjust_module_data(module)
    self.gather_parameters(module)
    for param in module.parameters(recurse=False):
        param.data = param.colo_attr.data_payload
        assert param.data.device.type == 'cuda', f"PRE BWD param.data must be on CUDA"