From 0f05015140467c80395308f98748c8cb85aab4b6 Mon Sep 17 00:00:00 2001 From: Marcin Ostrowski Date: Fri, 27 Sep 2024 17:58:19 -0700 Subject: [PATCH] Fix bug #2458 (#2459) Summary: https://github.com/pytorch/benchmark/issues/2458 Pull Request resolved: https://github.com/pytorch/benchmark/pull/2459 Reviewed By: xuzhao9 Differential Revision: D63476542 Pulled By: kit1980 fbshipit-source-id: 01e9db9cb03d34e82a773897417df2ccda410634 --- torchbenchmark/models/torch_multimodal_clip/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/torchbenchmark/models/torch_multimodal_clip/__init__.py b/torchbenchmark/models/torch_multimodal_clip/__init__.py index 6802ec932..fda18015f 100644 --- 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)]