device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 
torch.cuda.is_available() #True
word_model = word_model.to(device)
word_model.cuda()
for i,(x,y) in enumerate(dataset):
  if i < 1:
    x = x.cuda()
    output = word_model(x)
    i+=1
    print(output)
  if(i>=1):
    break

And it shows :
RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 'mat2' in call to _th_mm 
I guess that my input and model are not in the same device, but not sure why.