Function cast_training_params

src/diffusers/training_utils.py:139–146

(model: Union[torch.nn.Module, List[torch.nn.Module]], dtype=torch.float32)

def cast_training_params(model: Union[torch.nn.Module, List[torch.nn.Module]], dtype=torch.float32):
    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(