Bug in torch_multimodal_clip: IndexError: list index out of range

[ostrowskimarcin]

Hi, I was testing some of torchbench models and I found a simple bug that can be resolved.

Model: torch_multimodal_clip

Reproducing steps:

1. python3 install.py model torch_multimodal_clip
2. python3 test.py -v -k "test_torch_multimodal_clip_eval_cuda"

Error:

Running tests...

test_torch_multimodal_clip_eval_cuda (main.TestBenchmark) ... ERROR (9.481s)

======================================================================
ERROR [9.481s]: test_torch_multimodal_clip_eval_cuda (main.TestBenchmark)

torchbenchmark._components._impl.workers.subprocess_rpc.ChildTraceException: Traceback (most recent call last):
File "/root/torchbench/torchbenchmark/_components/_impl/workers/subprocess_rpc.py", line 482, in _run_block
exec( # noqa: P204
File "", line 2, in
File "/root/torchbench/torchbenchmark/util/model.py", line 317, in invoke
out = self.eval()
File "/root/torchbench/torchbenchmark/models/torch_multimodal_clip/init.py", line 91, in eval
return self.text[torch.argmax(score)]
IndexError: list index out of range

Solution:
The problem is inside init.py script. Default batch size is equal to 32, so the product of matrix multiplication returned as score is size [32,32]. In my opinion the goal is to find the highest score in each batch, but the method torch.argmax(score) returns global index (960). I think there's a missing argument of dim=1 to focus on each batch.
Moreover, after getting index or rather indices, in next step we want to get values from self.text field. I assume that we want to refer to self.texts according to proper batch_size.

I propose to change some lines in eval() method:

def eval(self):
        self.model.eval()
        image_embedding, text_embedding = self.model(
            self.image_tensor, self.text_tensor
        )
        score = image_embedding @ text_embedding.t()

        return self.text[torch.argmax(score)]

to

def eval(self):
        self.model.eval()
        image_embedding, text_embedding = self.model(
            self.image_tensor, self.text_tensor
        )
        score = image_embedding @ text_embedding.t()

        indices = torch.argmax(score, dim=1)
        return [self.texts[i][indices[i].item()] for i in range(self.batch_size)]

This works for any batch_size:

Running tests...

test_torch_multimodal_clip_eval_cuda (main.TestBenchmark) ... ok (9.999s)

---

Ran 1 test in 9.999s

OK