def _4D_bmm(self, batch1, batch2):
x =[]
for i in range(batch1.size(0)):
x.append( torch.bmm(torch.transpose(batch1[i], 1, 2), batch2[i]))
return Variable(torch.cat(x)).cuda()
And the exception raises form that model’s parameters are in GPU, while the input is in CPU.