Fix bug #2458 (#2459)

ostrowskimarcin · facebook-github-bot · commit 0f0501514046 · 2024-09-27T17:58:19.000-07:00
Summary: #2458 Pull Request resolved: #2459 Reviewed By: xuzhao9 Differential Revision: D63476542 Pulled By: kit1980 fbshipit-source-id: 01e9db9cb03d34e82a773897417df2ccda410634
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)]

Original file line number	Diff line number	Diff line change
`@@ -87,4 +87,5 @@ def eval(self):`
`87`	`87`	`)`
`88`	`88`	`score = image_embedding @ text_embedding.t()`
`89`	`89`
`90`		`- return self.text[torch.argmax(score)]`
	`90`	`+ indices = torch.argmax(score, dim=1)`
	`91`	`+ return [self.texts[i][indices[i].item()] for i in range(self.batch_size)]`