Similarity matrix shape does not match the shape of the mask

[hugofigueiras]

Hello,

I was currently testing the implementation when an error occured: The shape of the mask [512, 512] at index 0 does not match the shape of the indexed tensor [2, 2] at index 0.
My batch size is 256.

The error occurs in this part of the code:
similarity_matrix = torch.matmul(features, features.T)
mask = torch.eye(labels.shape[0], dtype=torch.bool).to(device)
labels = labels[~mask].view(labels.shape[0], -1)
similarity_matrix = similarity_matrix[~mask].view(similarity_matrix.shape[0], -1)

I'm wondering if this something I'm doing wrong and how do I match the shape of tensors?

Thanks in advance!

[laiyingxin2]

same question~

[sarda-devesh]

The issue is because of the following line of code

SimCLR/simclr.py

Line 28 in 1848fc9

labels = torch.cat([torch.arange(self.args.batch_size) for i in range(self.args.n_views)], dim=0)

which assumes that the number of features you have is a multiple of the batch size but this is not always true.

For example, consider a dataset with only a 100 elements and a batch size of 256. In that case it will create a labels of size (256, 256) even though it should only be (200, 200) (assuming you are using n_views of 2). The way to resolve this is by updating the above line to:

labels = torch.cat([torch.arange(int(features.size(0)/2)) for i in range(self.args.n_views)], dim=0)