Based on the error message it seems that main_out and Attentiom_map have shapes, which are incompatible for the torch.matmul operation.
In your code you’ve already added the print statements for the shapes, so you could (re-)use it and make sure both tensors have the expected shape.
Here is a small example of your error:
# works
a, b = torch.randn(2, 10, 20), torch.randn(2, 20, 10)
c = torch.matmul(a, b)

# raises your error
a, b = torch.randn(2, 10, 20), torch.randn(3, 20, 10)
c = torch.matmul(a, b)
> RuntimeError: The size of tensor a (2) must match the size of tensor b (3) at non-singleton dimension 0