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() |