I guess your training code might be missing the transferring of the data to the device via:
for data, target in loader:
    optimizer.zero_grad()
    data = data.to('cuda')
    target = target.to('cuda')
    output = model(data)
    ...

In your Dataset.__getitem__ method you can return CPUTensors.