diff options
author | Volpeon <git@volpeon.ink> | 2023-04-10 10:34:12 +0200 |
---|---|---|
committer | Volpeon <git@volpeon.ink> | 2023-04-10 10:34:12 +0200 |
commit | eb6a92abda5893c975437026cdaf0ce0bfefe2a4 (patch) | |
tree | a1525010b48362986e0cc2b7c3f7505a35dea71a /training/functional.py | |
parent | Update (diff) | |
download | textual-inversion-diff-eb6a92abda5893c975437026cdaf0ce0bfefe2a4.tar.gz textual-inversion-diff-eb6a92abda5893c975437026cdaf0ce0bfefe2a4.tar.bz2 textual-inversion-diff-eb6a92abda5893c975437026cdaf0ce0bfefe2a4.zip |
Update
Diffstat (limited to 'training/functional.py')
-rw-r--r-- | training/functional.py | 4 |
1 files changed, 1 insertions, 3 deletions
diff --git a/training/functional.py b/training/functional.py index e14aeea..46d25f6 100644 --- a/training/functional.py +++ b/training/functional.py | |||
@@ -644,11 +644,9 @@ def train( | |||
644 | min_snr_gamma: int = 5, | 644 | min_snr_gamma: int = 5, |
645 | **kwargs, | 645 | **kwargs, |
646 | ): | 646 | ): |
647 | text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler, extra = strategy.prepare( | 647 | text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler = strategy.prepare( |
648 | accelerator, text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler, **kwargs) | 648 | accelerator, text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler, **kwargs) |
649 | 649 | ||
650 | kwargs.update(extra) | ||
651 | |||
652 | vae.to(accelerator.device, dtype=dtype) | 650 | vae.to(accelerator.device, dtype=dtype) |
653 | vae.requires_grad_(False) | 651 | vae.requires_grad_(False) |
654 | vae.eval() | 652 | vae.eval() |