Fix bug #2458 (#2459)

Summary: #2458 Pull Request resolved: #2459 Reviewed By: xuzhao9 Differential Revision: D63476542 Pulled By: kit1980 fbshipit-source-id: 01e9db9cb03d34e82a773897417df2ccda410634
pytorch · Sep 28, 2024 · 0f05015 · 0f05015
1 parent 2edf80c
commit 0f05015
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/torchbenchmark/models/torch_multimodal_clip/__init__.py b/torchbenchmark/models/torch_multimodal_clip/__init__.py
@@ -87,4 +87,5 @@ def eval(self):
         )
         score = image_embedding @ text_embedding.t()
 
-        return self.text[torch.argmax(score)]
+        indices = torch.argmax(score, dim=1)
+        return [self.texts[i][indices[i].item()] for i in range(self.batch_size)]