def METHOD_NAME(self, module: torch.nn.Module):
    # record overall data statistics
    if self._memstarts_collector:
        self._memstarts_collector.sample_overall_data()
    for param in module.parameters(recurse=False):
        param.colo_attr.sharded_data_tensor.trans_state(TensorState.COMPUTE)
    # adjust stateful tensor to get enough CUDA memory
    self._stateful_tensor_mgr.adjust_layout()
    # record model data statistics
    if self._memstarts_collector:
        self._memstarts_collector.record_model_data_volume()