Update controllable_blender/generation_methods.py

controllable_blender/generation_methods.py

@@ -61,7 +61,7 @@ class RerankedTopKSampling(TreeSearch):
         
         if all_penalties.dim() == 1:
             all_penalties = all_penalties.unsqueeze(0)
-        all_penalties = all_penalties.expand(self.beam_size, -1)
+        all_penalties = all_penalties.repeat(self.beam_size, -1)
         
         penalties = torch.gather(all_penalties, -1, indices)
         penalised_probs = torch.mul(probs, penalties)