def unwrap_dist(model):
    """
    Unwrap the model from distributed training.
    """
    if isinstance(model, DDP):
        return model.module
    return model

    Unwrap the model from distributed training.
    
trellis/utils/dist_utils.py