def record_model_param_shape(model: torch.nn.Module) -> dict:
    """
    Record the shape of the model parameters.

    Args:
        model (torch.nn.Module): The model to record the parameter shape of.

    Returns:
        dict: The shape of the model parameters.
    """
    param_shape = {}
    for name, param in model.named_parameters():
        param_shape[name] = param.shape
    return param_shape

    Record the shape of the model parameters.

    Args:
        model (torch.nn.Module): The model to record the parameter shape of.

    Returns:
        dict: The shape of the model parameters.
    
opensora/utils/ckpt.py