(models, weights, output_dir)
| 400 | |
| 401 | # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format |
| 402 | def save_model_hook(models, weights, output_dir): |
| 403 | if accelerator.is_main_process: |
| 404 | for _, model in enumerate(models): |
| 405 | if isinstance(unwrap_model(model), VAEClass): |
| 406 | model.save_pretrained(os.path.join(output_dir, "vae")) |
| 407 | else: |
| 408 | raise ValueError(f"Unknown model type {model}") |
| 409 | |
| 410 | # make sure to pop weight so that corresponding model is not saved again |
| 411 | weights.pop() |
| 412 | |
| 413 | def load_model_hook(models, input_dir): |
| 414 | for _ in range(len(models)): |
nothing calls this directly
no test coverage detected