diff options
Diffstat (limited to 'training/functional.py')
| -rw-r--r-- | training/functional.py | 3 |
1 files changed, 3 insertions, 0 deletions
diff --git a/training/functional.py b/training/functional.py index 46d25f6..ff6d3a9 100644 --- a/training/functional.py +++ b/training/functional.py | |||
| @@ -695,5 +695,8 @@ def train( | |||
| 695 | callbacks=callbacks, | 695 | callbacks=callbacks, |
| 696 | ) | 696 | ) |
| 697 | 697 | ||
| 698 | accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=False) | ||
| 699 | accelerator.unwrap_model(unet, keep_fp32_wrapper=False) | ||
| 700 | |||
| 698 | accelerator.end_training() | 701 | accelerator.end_training() |
| 699 | accelerator.free_memory() | 702 | accelerator.free_memory() |