VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-09 06:45:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

Set grad enabled if is_train and train_text_encoder

Browse Source 
We only want to be enabling grad if we are training.
This commit is contained in:
rockerBOO
2023-11-07 11:59:30 -05:00
parent 9c591bdb12
commit 569ca72fc4
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
train_network.py
View File

@@ -145,7 +145,7 @@ class NetworkTrainer:
latents = latents * self.vae_scale_factor
b_size = latents.shape[0]
with torch.set_grad_enabled(train_text_encoder), accelerator.autocast():
with torch.set_grad_enabled(is_train and train_text_encoder), accelerator.autocast():
# Get the text embedding for conditioning
if args.weighted_captions:
text_encoder_conds = get_weighted_text_embeddings(