(models, weights, output_dir)
| 180 | if version.parse(accelerate.__version__) >= version.parse("0.16.0"): |
| 181 | # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format |
| 182 | def save_model_hook(models, weights, output_dir): |
| 183 | if cfgs.use_ema: |
| 184 | ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema")) |
| 185 | |
| 186 | for i, model in enumerate(models): |
| 187 | model.save_pretrained(os.path.join(output_dir, "unet")) |
| 188 | |
| 189 | # make sure to pop weight so that corresponding model is not saved again |
| 190 | weights.pop() |
| 191 | |
| 192 | def load_model_hook(models, input_dir): |
| 193 | if cfgs.use_ema: |
nothing calls this directly
no test coverage detected