tensor fix

controllable_blender/generation_methods.py

@@ -57,7 +57,12 @@ class RerankedTopKSampling(TreeSearch):
         values, indices = logprobs.topk(self.k, dim=-1)
         probs = torch.softmax(values, dim=-1)
 
-        all_penalties = self.reranker.token_penalties.repeat(self.beam_size, 1).to(probs.device)
+        all_penalties = self.reranker.token_penalties.to(probs.device)
+        
+        if all_penalties.dim() == 1:
+            all_penalties = all_penalties.unsqueeze(0)
+        all_penalties = all_penalties.expand(self.beam_size, -1)
+        
         penalties = torch.gather(all_penalties, -1, indices)
         penalised_probs = torch.mul(probs, penalties)