Repository: PennyLaneAI/pennylane

Issue:
Torchlayer error when running on GPU
This is a continuation to the post here: https://discuss.pennylane.ai/t/pennylane-and-pytorch-running-on-gpu/457. 
I have two separate codes – one with classical machine learning (nothing to do with pennylane) and one with quantum ml (below). The only difference between the two is that the qml has additional pennylane code in the DQN class. Running the classical ml code runs without a problem on the GPU but when I run the qml code I get an error. Here is part of the code:


```
import gym
import math
import random
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from collections import namedtuple
from itertools import count
from PIL import Image
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision.transforms as T
from torch.nn.functional import relu, sigmoid
import pennylane as qml
import time

out_dim = 2  # output dimension of model
wires = 1  # this is the width of the quantum element
n_quantum_layers = 2  # this is the depth of the quantum element

def layer(inputs, w0, w1, w2, w3, w4, w5, w6, w7, w8, w9, w10):
    qml.templates.SqueezingEmbedding(inputs, wires=range(wires))
    qml.templates.CVNeuralNetLayers(w0, w1, w2, w3, w4, w5, w6, w7, w8, w9, w10,
                                    wires=range(wires))
    return [qml.expval(qml.X(wires=i)) for i in range(wires)]


class DQN(nn.Module):

    def __init__(self, img_height, img_width):
        super().__init__()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(in_features=img_height * img_width * 3, out_features=12)
        self.fc2 = nn.Linear(in_features=12, out_features=8)
       # self.fc3 = nn.Linear(in_features=10, out_features=8)
        self.clayer_in = torch.nn.Linear(in_features=8, out_features=wires)
        self.clayer_out = torch.nn.Linear(wires, out_dim)

        dev = qml.device('strawberryfields.fock', wires=wires, cutoff_dim=3)
        self.layer_qnode = qml.QNode(layer, dev)

        weights = qml.init.cvqnn_layers_all(n_quantum_layers, wires)
        weight_shapes = {"w{}".format(i): w.shape for i, w in enumerate(weights)}
        
        self.qlayer = qml.qnn.TorchLayer(self.layer_qnode, weight_shapes)

    def forward(self, t):
        t = self.flatten(t)
        t = self.fc1(t)
        t = self.fc2(t)
       # t = self.fc3(t)
        t = self.clayer_in(t)
        t = self.qlayer(t)
        t = self.clayer_out(t)
        t = t.sigmoid()
        return t

batch_size = 128
gamma = 0.999
eps_start = 1
eps_end = 0.01
eps_decay = 0.0005
target_update = 10
memory_size = 500000
lr_start = 0.01
lr_end = 0.00001
lr_decay = 0.00009
num_episodes = 1000 # run for more episodes for better results

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

em = CartPoleEnvManager(device)
strategy = EpsilonGreedyStrategy(eps_start, eps_end, eps_decay)
agent = Agent(strategy, em.num_actions_available(), device)
memory = ReplayMemory(memory_size)
#learning_rate = LearningRate(lr_start,lr_end,lr_decay)
#learn = lr(learning_rate)

policy_net = DQN(em.get_screen_height(), em.get_screen_width()).to(device)
target_net = DQN(em.get_screen_height(), em.get_screen_width()).to(device)
target_net.load_state_dict(policy_net.state_dict())
target_net.eval() #tells pytorch that target_net is only used for inference, not training
optimizer = optim.Adam(params=policy_net.parameters(), lr=0.01)

i = 0
episode_durations = []
for episode in range(num_episodes): #iterate over each episode
    program_starts = time.time()
    em.reset()
    state = em.get_state()
    
    for timestep in count():
        action = agent.select_action(state, policy_net)
        reward = em.take_action(action)
        next_state = em.get_state()
        memory.push(Experience(state, action, next_state, reward))
        state = next_state
        #i+=1
        #print(i)
        if memory.can_provide_sample(batch_size):
            scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=100, gamma=0.9)
            experiences = memory.sample(batch_size)
            states, actions, rewards, next_states = extract_tensors(experiences)
            
            current_q_values = QValues.get_current(policy_net, states, actions)
            next_q_values = QValues.get_next(target_net, next_states) #will get the max qvalues of the next state, q values of next state are used via next state
            target_q_values = (next_q_values * gamma) + rewards

            loss = F.mse_loss(current_q_values, target_q_values.unsqueeze(1))
            optimizer.zero_grad() # sets the gradiesnt of all weights n biases in policy_net to zero
            loss.backward() #computes gradient of loss with respect to all weights n biases in the policy net
            optimizer.step() # updates the weights n biases with the gradients that were computed form loss.backwards
            scheduler.step()
        if em.done:
            episode_durations.append(timestep)
            plot(episode_durations, 100)
            break
    if episode % target_update == 0:
        target_net.load_state_dict(policy_net.state_dict()) 
    now = time.time()
    print("Episode hat {0} Sekunden gedauert".format(now - program_starts))     
        
em.close()
```

And when running the code, the following error appears:

```
Traceback (most recent call last):
  File "qdqn.py", line 328, in <module>
    loss.backward() #computes gradient of loss with respect to all weights n biases in the policy net
  File "/home/ubuntu/anaconda3/envs/gymm/lib/python3.8/site-packages/torch/tensor.py", line 198, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/home/ubuntu/anaconda3/envs/gymm/lib/python3.8/site-packages/torch/autograd/__init__.py", line 98, in backward
    Variable._execution_engine.run_backward(
RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 'mat2' in call to _th_mm
```



Discussion:
Thanks @mizkulg! This is definitely something that we will need to explore further; the `qnn.TorchLayer` class is a new feature, and has not yet been thoroughly tested on GPUs. Until this is fixed, we recommend using CPUs.
Thanks @mizkulg! Would you be able to provide a _minimal_ working example of the error? The code you shared above is great, but I imagine the error can be reproduced with much fewer lines of code. E.g., could we get this error with less than 10 lines? By doing this, it will be easier to characterize the error and zero in on how it can be fixed.
Hi, I have the same problem. Here is a reproducing code.

```
import pennylane as qml
import torch

n_qubits = 4
dev = qml.device("default.qubit", wires=n_qubits)

@qml.qnode(dev, interface='torch')
def circuit(inputs, weights):
    qml.templates.AmplitudeEmbedding(inputs, wires=range(n_qubits), pad=0.)
    qml.templates.BasicEntanglerLayers(weights, wires=range(n_qubits))
    return [qml.expval(qml.PauliZ(wires=i)) for i in range(n_qubits)]

n_layers = 1
weight_shapes = {"weights": (n_layers, n_qubits)}

qlayer = qml.qnn.TorchLayer(circuit, weight_shapes)

x = torch.rand((5, 10))
norm = x.norm(p=2, dim=1, keepdim=True)
x_normalized = x.div(norm).to(torch.device('cuda'))
loss = torch.sum(qlayer(x_normalized)).squeeze()
loss.backward()


```
Hello @shukob, thanks for your example! Could you please also provide the error message you get? 
@glassnotes here it is.
```
---------------------------------------------------------------------------

RuntimeError                              Traceback (most recent call last)

<ipython-input-1-ec0516067b63> in <module>
     21 x_normalized = x.div(norm).to(torch.device('cuda'))
     22 loss = torch.sum(qlayer(x_normalized)).squeeze()
---> 23 loss.backward()
     24 
     25 

~/anaconda3/envs/***/lib/python3.7/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
    219                 retain_graph=retain_graph,
    220                 create_graph=create_graph)
--> 221         torch.autograd.backward(self, gradient, retain_graph, create_graph)
    222 
    223     def register_hook(self, hook):

~/anaconda3/envs/***/lib/python3.7/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
    130     Variable._execution_engine.run_backward(
    131         tensors, grad_tensors_, retain_graph, create_graph,
--> 132         allow_unreachable=True)  # allow_unreachable flag
    133 
    134 

~/anaconda3/envs/***/lib/python3.7/site-packages/torch/autograd/function.py in apply(self, *args)
     87     def apply(self, *args):
     88         # _forward_cls is defined by derived class
---> 89         return self._forward_cls.backward(self, *args)  # type: ignore
     90 
     91 

~/anaconda3/envs/***/lib/python3.7/site-packages/torch/autograd/function.py in wrapper(ctx, *args)
    208     def wrapper(ctx, *args):
    209         with torch.no_grad():
--> 210             outputs = fn(ctx, *args)
    211 
    212         if not torch.is_grad_enabled():

~/anaconda3/envs/***/lib/python3.7/site-packages/pennylane/interfaces/torch.py in backward(ctx, grad_output)
    215             jacobian = torch.as_tensor(jacobian, dtype=grad_output.dtype)
    216 
--> 217             vjp = torch.transpose(grad_output.view(-1, 1), 0, 1) @ jacobian
    218             vjp = vjp.flatten()
    219 

RuntimeError: Tensor for argument #3 'mat2' is on CPU, but expected it to be on GPU (while checking arguments for addmm)
```

If I changed the line 215 in `interfaces/torch.py` into 
```
jacobian = torch.as_tensor(jacobian, dtype=grad_output.dtype).to(grad_output)
```
then another error occurrd.

```

---------------------------------------------------------------------------

RuntimeError                              Traceback (most recent call last)

<ipython-input-1-45fa4b49f4d4> in <module>
     20 x_normalized = x.div(norm).to(torch.device('cuda'))
     21 loss = torch.sum(qlayer(x_normalized)).squeeze()
---> 22 loss.backward()
     23 
     24 

~/anaconda3/envs/***/lib/python3.7/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
    219                 retain_graph=retain_graph,
    220                 create_graph=create_graph)
--> 221         torch.autograd.backward(self, gradient, retain_graph, create_graph)
    222 
    223     def register_hook(self, hook):

~/anaconda3/envs/***/lib/python3.7/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
    130     Variable._execution_engine.run_backward(
    131         tensors, grad_tensors_, retain_graph, create_graph,
--> 132         allow_unreachable=True)  # allow_unreachable flag
    133 
    134 

RuntimeError: Function _TorchQNodeBackward returned an invalid gradient at index 1 - expected type TensorOptions(dtype=float, device=cpu, layout=Strided, requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt)) but got TensorOptions(dtype=float, device=cuda:0, layout=Strided, requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))


```
Debugged `qnn/torch.py` & `interfaces/torch.py`  and found the problem is that the device for qnode params does not match `inputs `in `_evaluate_qnode`.

So the fix will be to change in `_evaluate_qnode` like below in addition to make `jacobian` device type mach `grad_output`
```
w = self.qnode_weights[arg].to(x)
```


Resolving patch:
diff --git a/pennylane/interfaces/torch.py b/pennylane/interfaces/torch.py
--- a/pennylane/interfaces/torch.py
+++ b/pennylane/interfaces/torch.py
@@ -212,7 +212,7 @@ def backward(ctx, grad_output):  # pragma: no cover
 
             # evaluate the Jacobian matrix of the QNode
             jacobian = qnode.jacobian(ctx.args, ctx.kwargs)
-            jacobian = torch.as_tensor(jacobian, dtype=grad_output.dtype)
+            jacobian = torch.as_tensor(jacobian, dtype=grad_output.dtype).to(grad_output)
 
             vjp = torch.transpose(grad_output.view(-1, 1), 0, 1) @ jacobian
             vjp = vjp.flatten()
diff --git a/pennylane/qnn/torch.py b/pennylane/qnn/torch.py
--- a/pennylane/qnn/torch.py
+++ b/pennylane/qnn/torch.py
@@ -330,7 +330,7 @@ def _evaluate_qnode(self, x):
 
         for arg in self.sig:
             if arg is not self.input_arg:  # Non-input arguments must always be positional
-                w = self.qnode_weights[arg]
+                w = self.qnode_weights[arg].to(x)
 
                 qnode = functools.partial(qnode, w)
             else:
@@ -349,7 +349,10 @@ def _evaluate_qnode_tape_mode(self, x):
         Returns:
             tensor: output datapoint
         """
-        kwargs = {**{self.input_arg: x}, **self.qnode_weights}
+        kwargs = {
+            **{self.input_arg: x},
+            **{arg: weight.to(x) for arg, weight in self.qnode_weights.items()},
+        }
         return self.qnode(**kwargs).type(x.dtype)
 
     def __str__(self):
diff --git a/pennylane/tape/interfaces/torch.py b/pennylane/tape/interfaces/torch.py
--- a/pennylane/tape/interfaces/torch.py
+++ b/pennylane/tape/interfaces/torch.py
@@ -73,7 +73,7 @@ def backward(ctx, grad_output):  # pragma: no cover
         jacobian = tape.jacobian(device, params=ctx.args, **tape.jacobian_options)
         tape.set_parameters(ctx.all_params, trainable_only=False)
 
-        jacobian = torch.as_tensor(jacobian, dtype=grad_output.dtype)
+        jacobian = torch.as_tensor(jacobian, dtype=grad_output.dtype).to(grad_output)
 
         vjp = grad_output.view(1, -1) @ jacobian
         grad_input_list = torch.unbind(vjp.flatten())


Test patch:
diff --git a/tests/qnn/test_qnn_torch.py b/tests/qnn/test_qnn_torch.py
--- a/tests/qnn/test_qnn_torch.py
+++ b/tests/qnn/test_qnn_torch.py
@@ -418,6 +418,33 @@ def test_gradients(self, get_circuit, n_qubits):
             assert torch.allclose(g1, g2)
         assert len(weights) == len(list(layer.parameters()))
 
+    @pytest.mark.parametrize("n_qubits, output_dim", indices_up_to(1))
+    def test_cuda_backward(self):
+        """Test if TorchLayer can be run on GPU"""
+        if not torch.cuda.is_available():
+            pytest.skip("Cuda device not available")
+        else:
+            try:
+                n_qubits = 4
+                dev = qml.device("default.qubit", wires=n_qubits)
+
+                @qml.qnode(dev, interface="torch")
+                def circuit(inputs, weights):
+                    qml.templates.AngleEmbedding(inputs, wires=range(n_qubits))
+                    qml.templates.BasicEntanglerLayers(weights, wires=range(n_qubits))
+                    return [qml.expval(qml.PauliZ(wires=i)) for i in range(n_qubits)]
+
+                n_layers = 1
+                weight_shapes = {"weights": (n_layers, n_qubits)}
+
+                qlayer = qml.qnn.TorchLayer(circuit, weight_shapes)
+
+                x = torch.rand((5, n_qubits), dtype=torch.float64)
+                loss = torch.sum(qlayer(x)).squeeze()
+                loss.backward()
+            except Exception:
+                pytest.fail("Exception raised in torch CUDA backward")
+
 
 @pytest.mark.parametrize("interface", ["autograd", "torch", "tf"])
 @pytest.mark.parametrize("n_qubits, output_dim", indices_up_to(1))