hotfix multiple prompts

shyamsn97 · Feb 16, 2023 · 70bbdb8 · 70bbdb8
1 parent d7c1ab7
commit 70bbdb8
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/mario_gpt/lm.py b/mario_gpt/lm.py
@@ -133,7 +133,11 @@ def sample(
         self.lm.eval()
         with torch.no_grad():
             if seed is None:
-                seed = self.tokenizer("X", return_tensors="pt").input_ids.view(1, 1)
+                seed = (
+                    self.tokenizer("X", return_tensors="pt")
+                    .input_ids.view(1, 1)
+                    .repeat(len(prompts), 1)
+                )
             out = seed.to(self.device)
             if encoder_hidden_states is None:
                 if prompts is not None: