TypeError: torch.baddbmm received an invalid combination of arguments - got (int, torch.FloatTensor, int, torch.FloatTensor, torch.cuda.FloatTensor, out=torch.FloatTensor), but expected one of:
 * (torch.FloatTensor source, torch.FloatTensor batch1, torch.FloatTensor batch2, *, torch.FloatTensor out)
 * (float beta, torch.FloatTensor source, torch.FloatTensor batch1, torch.FloatTensor batch2, *, torch.FloatTensor out)
 * (torch.FloatTensor source, float alpha, torch.FloatTensor batch1, torch.FloatTensor batch2, *, torch.FloatTensor out)
 * (float beta, torch.FloatTensor source, float alpha, torch.FloatTensor batch1, torch.FloatTensor batch2, *, torch.FloatTensor out)
      didn't match because some of the arguments have invalid types: (int, torch.FloatTensor, int, torch.FloatTensor, torch.cuda.FloatTensor, out=torch.FloatTensor)


The doc didn’t describe the argument should be a torch.cuda.FloatTensor or a torch.FloatTensor.
How can i know about it?
And how can i test it?
I’m very confused.