# RuntimeError with Mixed Device Tensors when Integrating PennyLane with PyTorch **URL:** \ **Category:** PennyLane Help\ **Created:** [March 20, 2024, 12:28pm UTC](https://discuss.pennylane.ai/t/runtimeerror-with-mixed-device-tensors-when-integrating-pennylane-with-pytorch/4318 "2024-03-20T12:28:03Z")\ **Posts on this page:** 4\ **Page:** 1
**Author:** ![linzhongyan1](https://yyz2.discourse-cdn.com/flex012/user_avatar/discuss.pennylane.ai/linzhongyan1/32/2432_2.png) [@linzhongyan1](https://discuss.pennylane.ai/u/linzhongyan1)\ **Post date:** [March 20, 2024, 12:28pm UTC](https://discuss.pennylane.ai/t/runtimeerror-with-mixed-device-tensors-when-integrating-pennylane-with-pytorch/4318/1 "2024-03-20T12:28:03Z")
I’m encountering a runtime error while trying to integrate PennyLane with a PyTorch model for a hybrid quantum-classical neural network project. The error occurs during the forward pass of my model, specifically when executing a quantum circuit defined with PennyLane and integrating its output with PyTorch tensors. Here’s the error message: vbnetCopy code ```auto RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu, cuda:0! ``` This error arises when I attempt to pass tensors from my PyTorch model, which is on a GPU, to a PennyLane quantum node, which operates on the CPU. I’ve made sure to move all relevant tensors to the CPU before executing the quantum node and back to the GPU afterwards, yet the error persists. Below is a simplified version of the code that leads to this error: pythonCopy code ```auto import torch import pennylane as qml from torch import nn from pennylane import numpy as np device = torch.device("cuda" if torch.cuda.is_available() else "cpu") #%% md ## Load Dataset There are two columns in the dataset: label and text. The column "text" holds the message content, whereas "label" is a binary variable where 1 indicates that the message is spam and 0 indicates that it is not spam. #%% df = pd.read_csv("data/spamdata_v2.csv") df.head() #%% df.shape #%% # check class distribution df['label'].value_counts(normalize = True) #%% #%% md This dataset will now be divided into three sets: train, validation, and test. We will fine-tune the model using the train set and the validation set, and make predictions for the test set. #%% train_text, temp_text, train_labels, temp_labels = train_test_split(df['text'], df['label'], random_state=2018, test_size=0.3, stratify=df['label']) # we will use temp_text and temp_labels to create validation and test set val_text, test_text, val_labels, test_labels = train_test_split(temp_text, temp_labels, random_state=2018, test_size=0.5, stratify=temp_labels) #%% md ## Import BERT Model and BERT Tokenizer #%% # import BERT-base pretrained model bert = AutoModel.from_pretrained('bert-base-uncased') # Load the BERT tokenizer tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased') #%% md ## Tokenize the Sentences Since the messages (text) in the dataset are of variable lengths, we will employ padding to ensure that the length of each message is the same. We can pad messages using the maximum sequence length. To determine the correct padding length, we may also examine the distribution of sequence lengths in the train set. #%% # get length of all the messages in the train set seq_len = [len(i.split()) for i in train_text] pd.Series(seq_len).hist(bins = 30) #%% md It is evident that the majority of the texts include little more than 25 words. In contrast to the maximum length of 175, If we choose 175 as the padding length, then all input sequences will have a length of 175 and the majority of tokens in those sequences will be padding tokens, which will not help the model learn anything useful and will also slow down training. We will thus set the padding length to 25. #%% max_seq_len = 25 #%% # tokenize and encode sequences in the training set tokens_train = tokenizer.batch_encode_plus( train_text.tolist(), max_length = max_seq_len, pad_to_max_length=True, truncation=True, return_token_type_ids=False ) # tokenize and encode sequences in the validation set tokens_val = tokenizer.batch_encode_plus( val_text.tolist(), max_length = max_seq_len, pad_to_max_length=True, truncation=True, return_token_type_ids=False ) # tokenize and encode sequences in the test set tokens_test = tokenizer.batch_encode_plus( test_text.tolist(), max_length = max_seq_len, pad_to_max_length=True, truncation=True, return_token_type_ids=False ) #%% md ## convert the integer sequences to tensors. #%% # for train set train_seq = torch.tensor(tokens_train['input_ids']) train_mask = torch.tensor(tokens_train['attention_mask']) train_y = torch.tensor(train_labels.tolist()) # for validation set val_seq = torch.tensor(tokens_val['input_ids']) val_mask = torch.tensor(tokens_val['attention_mask']) val_y = torch.tensor(val_labels.tolist()) # for test set test_seq = torch.tensor(tokens_test['input_ids']) test_mask = torch.tensor(tokens_test['attention_mask']) test_y = torch.tensor(test_labels.tolist()) #%% md ## Create DataLoaders Now, dataloaders will be created for both the train and validation sets. During the training phase, these dataloaders will send batches of train data and validation data to the model as input. #%% from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler #define a batch size batch_size = 32 # wrap tensors train_data = TensorDataset(train_seq, train_mask, train_y) # sampler for sampling the data during training train_sampler = RandomSampler(train_data) # dataLoader for train set train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=batch_size) # wrap tensors val_data = TensorDataset(val_seq, val_mask, val_y) # sampler for sampling the data during training val_sampler = SequentialSampler(val_data) # dataLoader for validation set val_dataloader = DataLoader(val_data, sampler = val_sampler, batch_size=batch_size) #%% md ## Freeze BERT Parameters We freeze all the layers of the BERT model and attach a few neural network layers of our own and train this new model. Note that the weights of only the attached layers will be updated during model training. #%% # freeze all the parameters for param in bert.parameters(): param.requires_grad = False #%% md ## Define Model Architecture Hybrid transfer learning model (classical-to-quantum). Setting of the main parameters of the network model and of the training process. #%% n_qubits = 5 # Number of qubits q_depth = 4 # Depth of the quantum circuit (number of variational layers) max_layers = 15 # Keep 15 even if not all are used. q_delta = 0.01 # Initial spread of random quantum weights #%% md Let us initialize a PennyLane device with the default simulator. #%% dev = qml.device('default.qubit', wires=n_qubits) class Quantumnet(nn.Module): def __init__ (self, bert): super(Quantumnet, self). __init__ () self.bert = bert self.pre_net = nn.Linear(768, n_qubits) self.q_params = nn.Parameter(q_delta * torch.randn(max_layers * n_qubits)) self.post_net = nn.Linear(n_qubits, 2) def forward(self, sent_id, mask): _, cls_hs = self.bert(sent_id, attention_mask=mask, return_dict=False) pre_out = self.pre_net(cls_hs) q_in = torch.tanh(pre_out) * np.pi / 2.0 # Apply the quantum circuit to each element of the batch and append to q_out q_out = torch.Tensor(0, n_qubits) q_out = q_out.to(device) for elem in q_in: q_out_elem = q_net(elem,self.q_params).float().unsqueeze(0) q_out = torch.cat((q_out, q_out_elem)) return self.post_net(q_out) dev = qml.device('default.qubit', wires=n_qubits) @qml.qnode(dev, interface='torch') def q_net(q_in, q_weights_flat): # Reshape weights q_weights = q_weights_flat.reshape(max_layers, n_qubits) # Start from state |+> , unbiased w.r.t. |0> and |1> H_layer(n_qubits) # Embed features in the quantum node RY_layer(q_in) # Sequence of trainable variational layers for k in range(q_depth): entangling_layer(n_qubits) RY_layer(q_weights[k + 1]) # Expectation values in the Z basis return [qml.expval(qml.PauliZ(j)) for j in range(n_qubits)] model = Quantumnet(bert, n_qubits, max_layers, q_delta).to(device) # push the model to GPU model = model.to(device) from transformers import AdamW # define the optimizer optimizer = AdamW(model.parameters(), lr = 1e-3) #%% md #%% md ## Find Class Weights Our dataset has an imbalance between classes. The vast majority of messages are not spam. Therefore, we will first calculate class weights for the labels in the train set and then send these weights to the loss function so that the class imbalance is taken care of. #%% from sklearn.utils.class_weight import compute_class_weight #compute the class weights class_wts = compute_class_weight(class_weight = "balanced", classes = np.unique(train_labels), y = train_labels) print(class_wts) #%% # convert class weights to tensor weights= torch.tensor(class_wts,dtype=torch.float) weights = weights.to(device) # loss function cross_entropy = nn.CrossEntropyLoss(weight=weights) # number of training epochs epochs = 10 model = model.to(device) #%% md ## Fine-Tune BERT So far, we have described the model's architecture, specified the optimizer and loss function, and prepared the dataloaders. Now we must define a couple of functions to train (fine-tune) and evaluate the model, respectively. #%% # function to train the model def train(): model.train() total_loss, total_accuracy = 0, 0 # empty list to save model predictions total_preds=[] # iterate over batches for step,batch in enumerate(train_dataloader): # progress update after every 50 batches. if step % 50 == 0 and not step == 0: print(' Batch {:>5,} of {:>5,}.'.format(step, len(train_dataloader))) # push the batch to gpu batch = [r.to(device) for r in batch] sent_id, mask, labels = batch # clear previously calculated gradients model.zero_grad() # get model predictions for the current batch preds = model(sent_id, mask) # compute the loss between actual and predicted values loss = cross_entropy(preds, labels) # add on to the total loss total_loss = total_loss + loss.item() # backward pass to calculate the gradients loss.backward() # clip the the gradients to 1.0. It helps in preventing the exploding gradient problem torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # update parameters optimizer.step() # model predictions are stored on GPU. So, push it to CPU preds=preds.detach().cpu().numpy() # append the model predictions total_preds.append(preds) # compute the training loss of the epoch avg_loss = total_loss / len(train_dataloader) # predictions are in the form of (no. of batches, size of batch, no. of classes). # reshape the predictions in form of (number of samples, no. of classes) total_preds = np.concatenate(total_preds, axis=0) #returns the loss and predictions return avg_loss, total_preds #%% # function for evaluating the model def evaluate(): print("\nEvaluating...") # deactivate dropout layers model.eval() total_loss, total_accuracy = 0, 0 # empty list to save the model predictions total_preds = [] # iterate over batches for step,batch in enumerate(val_dataloader): # Progress update every 50 batches. if step % 50 == 0 and not step == 0: # Calculate elapsed time in minutes. elapsed = format_time(time.time() - t0) # Report progress. print(' Batch {:>5,} of {:>5,}.'.format(step, len(val_dataloader))) # push the batch to gpu batch = [t.to(device) for t in batch] sent_id, mask, labels = batch # deactivate autograd with torch.no_grad(): # model predictions preds = model(sent_id, mask) # compute the validation loss between actual and predicted values loss = cross_entropy(preds,labels) total_loss = total_loss + loss.item() preds = preds.detach().cpu().numpy() total_preds.append(preds) # compute the validation loss of the epoch avg_loss = total_loss / len(val_dataloader) # reshape the predictions in form of (number of samples, no. of classes) total_preds = np.concatenate(total_preds, axis=0) return avg_loss, total_preds #%% md ## Start Model Training #%% # set initial loss to infinite best_valid_loss = float('inf') weights= torch.tensor(class_wts,dtype=torch.float) weights = weights.to(device) # loss function print(class_wts) # empty lists to store training and validation loss of each epoch train_losses=[] valid_losses=[] #for each epoch for epoch in range(epochs): print('\n Epoch {:} / {:}'.format(epoch + 1, epochs)) #train model train_loss, _ = train() #evaluate model valid_loss, _ = evaluate() #save the best model if valid_loss < best_valid_loss: best_valid_loss = valid_loss torch.save(model.state_dict(), 'saved_weights.pt') # append training and validation loss train_losses.append(train_loss) valid_losses.append(valid_loss) print(f'\nTraining Loss: {train_loss:.3f}') print(f'Validation Loss: {valid_loss:.3f}') I've ensured that all tensors are correctly moved to the CPU before calling the quantum node and then moved back to the GPU as necessary. Despite this, the error suggests there's still a mismatch in device allocation. Could you please help me understand what might be causing this issue and how to resolve it? Am I missing a step in correctly handling device allocation between PennyLane and PyTorch tensors? Thank you for your assistance! pytorch 2.2.0 pennylane-0.35.1 pennylane-lightning-0.35.1 cuda 4060 . ``` ---
**Author:** ![isaacdevlugt](https://yyz2.discourse-cdn.com/flex012/user_avatar/discuss.pennylane.ai/isaacdevlugt/32/1159_2.png) [@isaacdevlugt](https://discuss.pennylane.ai/u/isaacdevlugt)\ **Post date:** [March 20, 2024, 10:51pm UTC](https://discuss.pennylane.ai/t/runtimeerror-with-mixed-device-tensors-when-integrating-pennylane-with-pytorch/4318/2 "2024-03-20T22:51:00Z")
Hey @linzhongyan1, welcome to the forum! 🕶 I think I’m missing the definition of `bert` in your code. Can you post your full code for me to copy-paste on my side and try to replicate your problem? ---
**Author:** ![linzhongyan1](https://yyz2.discourse-cdn.com/flex012/user_avatar/discuss.pennylane.ai/linzhongyan1/32/2432_2.png) [@linzhongyan1](https://discuss.pennylane.ai/u/linzhongyan1)\ **Post date:** [March 21, 2024, 4:37am UTC](https://discuss.pennylane.ai/t/runtimeerror-with-mixed-device-tensors-when-integrating-pennylane-with-pytorch/4318/4 "2024-03-21T04:37:24Z")
## Thank you for your reply. I have posted full code ，After running this segment of code, the following error occurred: RuntimeError Traceback (most recent call last) Cell In[464], line 23 20 print(‘\n Epoch {:} / {:}’.format(epoch + 1, epochs)) 22 #train model —\> 23 train\_loss, \_ = train() 25 #evaluate model 26 valid\_loss, \_ = evaluate() Cell In[462], line 28, in train() 25 model.zero\_grad() 27 # get model predictions for the current batch —\> 28 preds = model(sent\_id, mask) 30 # compute the loss between actual and predicted values 31 loss = cross\_entropy(preds, labels) File ~.conda\envs\pytorch\Lib\site-packages\torch\nn\modules\module.py:1511, in Module.\_wrapped\_call\_impl(self, \*args, \*\*kwargs) 1509 return self.\_compiled\_call\_impl(\*args, \*\*kwargs) # type: ignore[misc] 1510 else: → 1511 return self.\_call\_impl(\*args, \*\*kwargs) File ~.conda\envs\pytorch\Lib\site-packages\torch\nn\modules\module.py:1520, in Module.\_call\_impl(self, \*args, \*\*kwargs) 1515 # If we don’t have any hooks, we want to skip the rest of the logic in 1516 # this function, and just call forward. 1517 if not (self.\_backward\_hooks or self.\_backward\_pre\_hooks or self.\_forward\_hooks or self.\_forward\_pre\_hooks 1518 or \_global\_backward\_pre\_hooks or \_global\_backward\_hooks 1519 or \_global\_forward\_hooks or \_global\_forward\_pre\_hooks): → 1520 return forward\_call(\*args, \*\*kwargs) 1522 try: 1523 result = None Cell In[457], line 25, in Quantumnet.forward(self, sent\_id, mask) 23 # 量子层操作在CPU上执行 24 for elem in q\_in.to(‘cpu’): —\> 25 q\_out\_elem = q\_net(elem, self.q\_params).float().unsqueeze(0) # q\_net的定义应该处理在CPU上的数据 26 q\_out = torch.cat((q\_out, q\_out\_elem), dim=0) 28 # 将量子层的输出移回模型确定的设备 File ~.conda\envs\pytorch\Lib\site-packages\pennylane\qnode.py:1039, in QNode. **call** (self, \*args, \*\*kwargs) 1034 full\_transform\_program.\_set\_all\_argnums( 1035 self, args, kwargs, argnums 1036 ) # pylint: disable=protected-access 1038 # pylint: disable=unexpected-keyword-arg → 1039 res = qml.execute( 1040 (self.\_tape,), 1041 device=self.device, 1042 gradient\_fn=self.gradient\_fn, 1043 interface=self.interface, 1044 transform\_program=full\_transform\_program, 1045 config=config, 1046 gradient\_kwargs=self.gradient\_kwargs, 1047 override\_shots=override\_shots, 1048 \*\*self.execute\_kwargs, 1049 ) 1051 res = res[0] 1053 # convert result to the interface in case the qfunc has no parameters File ~.conda\envs\pytorch\Lib\site-packages\pennylane\interfaces\execution.py:648, in execute(tapes, device, gradient\_fn, interface, transform\_program, config, grad\_on\_execution, gradient\_kwargs, cache, cachesize, max\_diff, override\_shots, expand\_fn, max\_expansion, device\_batch\_transform, device\_vjp) 646 # Exiting early if we do not need to deal with an interface boundary 647 if no\_interface\_boundary\_required: → 648 results = inner\_execute(tapes) 649 return post\_processing(results) 651 \_grad\_on\_execution = False File ~.conda\envs\pytorch\Lib\site-packages\pennylane\interfaces\execution.py:261, in \*make\_inner\_execute..inner\_execute(tapes, \*\*\*) 259 if numpy\_only: 260 tapes = tuple(qml.transforms.convert\_to\_numpy\_parameters(t) for t in tapes) → 261 return cached\_device\_execution(tapes) File ~.conda\envs\pytorch\Lib\site-packages\pennylane\interfaces\execution.py:383, in cache\_execute..wrapper(tapes, \*\*kwargs) 378 return (res, ) if return\_tuple else res 380 else: 381 # execute all unique tapes that do not exist in the cache 382 # convert to list as new device interface returns a tuple → 383 res = list(fn(tuple(execution\_tapes.values()), \*\*kwargs)) 385 final\_res = 387 for i, tape in enumerate(tapes): File ~.conda\envs\pytorch\Lib\site-packages\pennylane\devices\default\_qubit.py:523, in DefaultQubit.execute(self, circuits, execution\_config) 517 interface = ( 518 execution\_config.interface 519 if execution\_config.gradient\_method in {“backprop”, None} 520 else None 521 ) 522 if max\_workers is None: → 523 results = tuple( 524 simulate( 525 c, 526 rng=self.\_rng, 527 prng\_key=self.\_prng\_key, 528 debugger=self.\_debugger, 529 interface=interface, 530 state\_cache=self.\_state\_cache, 531 ) 532 for c in circuits 533 ) 534 else: 535 vanilla\_circuits = [convert\_to\_numpy\_parameters(c) for c in circuits] File ~.conda\envs\pytorch\Lib\site-packages\pennylane\devices\default\_qubit.py:524, in (.0) 517 interface = ( 518 execution\_config.interface 519 if execution\_config.gradient\_method in {“backprop”, None} 520 else None 521 ) 522 if max\_workers is None: 523 results = tuple( → 524 simulate( 525 c, 526 rng=self.\_rng, 527 prng\_key=self.\_prng\_key, 528 debugger=self.\_debugger, 529 interface=interface, 530 state\_cache=self.\_state\_cache, 531 ) 532 for c in circuits 533 ) 534 else: 535 vanilla\_circuits = [convert\_to\_numpy\_parameters(c) for c in circuits] File ~.conda\envs\pytorch\Lib\site-packages\pennylane\devices\qubit\simulate.py:240, in simulate(circuit, rng, prng\_key, debugger, interface, state\_cache) 203 def simulate( 204 circuit: qml.tape.QuantumScript, 205 rng=None, (…) 209 state\_cache: Optional[dict] = None, 210 ) → Result: 211 “”“Simulate a single quantum script. 212 213 This is an internal function that will be called by the successor to `default.qubit`. (…) 238 239 “”” → 240 state, is\_state\_batched = get\_final\_state(circuit, debugger=debugger, interface=interface) 241 if state\_cache is not None: 242 state\_cache[circuit.hash] = state File ~.conda\envs\pytorch\Lib\site-packages\pennylane\devices\qubit\simulate.py:128, in get\_final\_state(circuit, debugger, interface) 126 is\_state\_batched = bool(prep and prep.batch\_size is not None) 127 for op in circuit.operations[bool(prep) :]: → 128 state = apply\_operation(op, state, is\_state\_batched=is\_state\_batched, debugger=debugger) 130 # Handle postselection on mid-circuit measurements 131 if isinstance(op, qml.Projector): File ~.conda\envs\pytorch\Lib\functools.py:909, in singledispatch..wrapper(\*args, \*\*kw) 905 if not args: 906 raise TypeError(f’{funcname} requires at least ’ 907 ‘1 positional argument’) → 909 return dispatch(args[0]. **class** )(\*args, \*\*kw) File ~.conda\envs\pytorch\Lib\site-packages\pennylane\devices\qubit\apply\_operation.py:198, in apply\_operation(op, state, is\_state\_batched, debugger) 150 @singledispatch 151 def apply\_operation( 152 op: qml.operation.Operator, state, is\_state\_batched: bool = False, debugger=None 153 ): 154 “”“Apply and operator to a given state. 155 156 Args: (…) 196 197 “”” → 198 return \_apply\_operation\_default(op, state, is\_state\_batched, debugger) File ~.conda\envs\pytorch\Lib\site-packages\pennylane\devices\qubit\apply\_operation.py:208, in \_apply\_operation\_default(op, state, is\_state\_batched, debugger) 202 “”“The default behaviour of apply\_operation, accessed through the standard dispatch 203 of apply\_operation, as well as conditionally in other dispatches.”“” 204 if ( 205 len(op.wires) \< EINSUM\_OP\_WIRECOUNT\_PERF\_THRESHOLD 206 and math.ndim(state) \< EINSUM\_STATE\_WIRECOUNT\_PERF\_THRESHOLD 207 ) or (op.batch\_size and is\_state\_batched): → 208 return apply\_operation\_einsum(op, state, is\_state\_batched=is\_state\_batched) 209 return apply\_operation\_tensordot(op, state, is\_state\_batched=is\_state\_batched) File ~.conda\envs\pytorch\Lib\site-packages\pennylane\devices\qubit\apply\_operation.py:100, in apply\_operation\_einsum(op, state, is\_state\_batched) 97 op.\_batch\_size = batch\_size # pylint:disable=protected-access 98 reshaped\_mat = math.reshape(mat, new\_mat\_shape) → 100 return math.einsum(einsum\_indices, reshaped\_mat, state) File ~.conda\envs\pytorch\Lib\site-packages\pennylane\math\multi\_dispatch.py:539, in einsum(indices, like, optimize, \*operands) 537 if like is None: 538 like = get\_interface(\*operands) → 539 operands = np.coerce(operands, like=like) 540 if optimize is None or like == “torch”: 541 # torch einsum doesn’t support the optimize keyword argument 542 return np.einsum(indices, \*operands, like=like) File ~.conda\envs\pytorch\Lib\site-packages\autoray\autoray.py:80, in do(fn, like, \*args, \*\*kwargs) 31 “”“Do function named `fn` on `(*args, **kwargs)`, peforming single 32 dispatch to retrieve `fn` based on whichever library defines the class of 33 the `args[0]`, or the `like` keyword argument if specified. (…) 77 \ 78 “”” 79 backend = choose\_backend(fn, \*args, like=like, \*\*kwargs) —\> 80 return get\_lib\_fn(backend, fn)(\*args, \*\*kwargs) File ~.conda\envs\pytorch\Lib\site-packages\pennylane\math\single\_dispatch.py:601, in \_coerce\_types\_torch(tensors) 598 if len(device\_set) \> 1: # pragma: no cover 599 # GPU specific case 600 device\_names = ", “.join(str(d) for d in device\_set) → 601 raise RuntimeError( 602 f"Expected all tensors to be on the same device, but found at least two devices, {device\_names}!” 603 ) 605 device = device\_set.pop() if len(device\_set) == 1 else None 606 tensors = [torch.as\_tensor(t, device=device) for t in tensors] RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu, cuda:0! ---
**Author:** ![isaacdevlugt](https://yyz2.discourse-cdn.com/flex012/user_avatar/discuss.pennylane.ai/isaacdevlugt/32/1159_2.png) [@isaacdevlugt](https://discuss.pennylane.ai/u/isaacdevlugt)\ **Post date:** [March 26, 2024, 1:46pm UTC](https://discuss.pennylane.ai/t/runtimeerror-with-mixed-device-tensors-when-integrating-pennylane-with-pytorch/4318/5 "2024-03-26T13:46:46Z")
Hey @linzhongyan1, It’s a little hard to tell what’s going on because your code base is quite large. That, and I can’t run your code because I’m missing your input data. That said, it might be that your training data is the culprit; it might be on the cpu and everything else is on the gpu. I think if you do `torch.set_default_tensor_type('torch.cuda.FloatTensor')`, then that makes all torch tensors created be on the gpu (see here: [How to create a tensor on GPU as default - PyTorch Forums](https://discuss.pytorch.org/t/how-to-create-a-tensor-on-gpu-as-default/2128)). Let me know if that helps!