fix for #6700

3 years ago · 11183b4d90
parent e452facef4
commit 11183b4d90
1 changed files with 1 additions and 1 deletions
--- a/modules/textual_inversion/dataset.py
+++ b/modules/textual_inversion/dataset.py
@ -115,7 +115,7 @@ class PersonalizedBase(Dataset):
                weight /= weight.mean()
            elif use_weight:
                #If an image does not have a alpha channel, add a ones weight map anyway so we can stack it later
-                weight = torch.ones([channels] + latent_size)
+                weight = torch.ones(latent_sample.shape)
            else:
                weight = None