criterion = nn.MSELoss()
criterion(a, b)

This is dtype=torch.float for a and dtype=torch.int64 for b
So, both change to float