Thanks for the code. It looks alright, so could you add these print statements into the training loop:
inputs, labels = inputs.to(device), labels.to(device)
print(type(inputs))
print(type(labels))
print(inputs.size())
print(inputs.dim())