Quantum Transfer Learning

PennyLane Help
8

Hey @kevinkawchak,

It looks like somewhere along the line you’re creating a torch tensor on a cpu device instead of a gpu device (or the other way around). You’ll need to do something like this:

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
input_tensor = input_tensor.to(device)
output_tensor = output_tensor.to(device)

or you can set the default device like this (see here: torch.set_default_device — PyTorch 2.1 documentation)

>>> torch.tensor([1.2, 3]).device
device(type='cpu')
>>> torch.set_default_device('cuda')  # current device is 0
>>> torch.tensor([1.2, 3]).device
device(type='cuda', index=0)
>>> torch.set_default_device('cuda:1')
>>> torch.tensor([1.2, 3]).device
device(type='cuda', index=1)

Let me know if this helps!

1 Like