(model: Union[torch.nn.Module, List[torch.nn.Module]], dtype=torch.float32)
| 194 | |
| 195 | |
| 196 | def cast_training_params(model: Union[torch.nn.Module, List[torch.nn.Module]], dtype=torch.float32): |
| 197 | if not isinstance(model, list): |
| 198 | model = [model] |
| 199 | for m in model: |
| 200 | for param in m.parameters(): |
| 201 | # only upcast trainable parameters into fp32 |
| 202 | if param.requires_grad: |
| 203 | param.data = param.to(dtype) |
| 204 | |
| 205 | |
| 206 | def _set_state_dict_into_text_encoder( |
no test coverage detected