test tensor adjustment

controllable_blender/generation_methods.py

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