Inefficient implementation

model.py

@@ -52,17 +52,13 @@ class ImageCaptionModel:
 		
 		"""
 		pixel_values = self.processor(images=image, return_tensors="pt").pixel_values.to(self.device)
-			
-		if diversity_penalty != 0.0:
-			num_beam_groups = 2
-			num_captions = num_captions if num_captions % 2 == 0 else num_captions + 1
 
 		generated_ids = self.model.generate(
 			pixel_values=pixel_values,
 			max_length=max_length,
-			num_beams=num_captions,
+			num_beams=num_captions*2,
 			num_beam_groups=num_beam_groups,
-			num_return_sequences=num_captions,
+			num_return_sequences=num_captions*2,
 			temperature=temperature,
 			top_k=top_k,
 			top_p=top_p,
@@ -71,8 +67,9 @@ class ImageCaptionModel:
 		)
 
 		generated_caption = self.processor.batch_decode(generated_ids, skip_special_tokens=True)
+		generated_caption = [generated_caption[i] for i in range(0, num_captions*2, 2)]
 
-		return generated_caption[:num_captions]
+		return generated_caption
 
 
 class GitBaseCocoModel(ImageCaptionModel):