Function cast_training_params

src/diffusers/training_utils.py:316–330

Casts the training parameters of the model to the specified data type. Args: model: The PyTorch model whose parameters will be cast. dtype: The data type to which the model parameters will be cast.

(model: torch.nn.Module | list[torch.nn.Module], dtype=torch.float32)

def cast_training_params(model: torch.nn.Module | list[torch.nn.Module], dtype=torch.float32):
    """
    Casts the training parameters of the model to the specified data type.

    Args:
        model: The PyTorch model whose parameters will be cast.
        dtype: The data type to which the model parameters will be cast.
    """
    if not isinstance(model, list):
        model = [model]
    for m in model:
        for param in m.parameters():
            # only upcast trainable parameters into fp32
            if param.requires_grad:
                param.data = param.to(dtype)


def _set_state_dict_into_text_encoder(