forgot to normalise

ShoukanLabs · Oct 9, 2023 · 1574763 · 1574763
1 parent 1c129e7
commit 1574763
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/muse_maskgit_pytorch/trainers/maskgit_trainer.py b/muse_maskgit_pytorch/trainers/maskgit_trainer.py
@@ -171,13 +171,12 @@ def train(self):
                                 input_ids, attn_mask, self.model.transformer.t5, self.accelerator.device
                             )
                 else:
-                    img_for_embed = transforms.ToPILImage(imgs)
-
                     model, _, preprocess = self.clip_model
                     text = open_clip.tokenize(text)
 
                     with torch.no_grad():
                         text_embeds = model.encode_text(text)
+                        text_embeds /= text_embeds.norm(dim=-1, keepdim=True)
 
                 with self.accelerator.accumulate(self.model), self.accelerator.autocast():
                     loss = self.model(imgs, text_embeds=text_embeds)