input1 = torch.randn(1, 3, 10, 8) 
input2 = torch.randn(1, 1, 128, 8)
cos = nn.CosineSimilarity(dim=2, eps=1e-6)
output = cos(input1, input2)

Expected output shape is (1, 3, 10*128, 1) =>  (1, 3, 1280, 1)