Tensor.cpu() not working for me

Kaiwen_Xie: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/e/de8d468f2c94b122bf5203f2466c36993212a4ae.png" data-download-href="https://discuss.pytorch.org/uploads/default/de8d468f2c94b122bf5203f2466c36993212a4ae" title="Screenshot 2024-03-11 180658"><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/e/de8d468f2c94b122bf5203f2466c36993212a4ae.png" alt="Screenshot 2024-03-11 180658" data-base62-sha1="vKMGaalt4CIagGMBTMJtwVZLLhs" width="610" height="500" data-dominant-color="393939"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2024-03-11 180658</span><span class="informations">883×723 33.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>y_preds is a tensor thats on the GPU, but since Matplotlib can only work with NumPy, i tried to move it to the CPU by using the .cpu() method but the same error shows up.</p>
<p>I also tried the .cpu().numpy(), but the same error shows up. Could anyone advice me on how to solve this?</p>

ptrblck: <p>Could you post a minimal and executable code snippet reproducing the error, please?</p>

Kaiwen_Xie: <pre><code class="lang-auto">import torch
from torch import nn
import matplotlib.pyplot as plt
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

#Create some data using the linear regression formula

def plot_prediction(train_data = X_train,
                    train_labels = y_train,
                    test_data = X_test,
                    test_labels = y_test,
                    predictions = None):

  plt.figure(figsize=(10, 7))
  plt.scatter(train_data, train_labels, c="b", s = 4, label = "Training data")
  plt.scatter(test_data, test_labels, c="g", s = 4, label = "Test data")

  if predictions is not None:
    plt.scatter(X_test, predictions, c="r", label = "predictions")

  plt.legend(prop={"size" : 14})

weight = 0.7
bias = 0.3

#Create range values
start = 0
end = 1
step = 0.02

#Create X and y (features and labels)
X = torch.arange(start, end, step).unsqueeze(dim = 1) #without unsqueeze, errors will pop up, shape errors during matmul in training loop
y = weight * X + bias
#.unsqueeze(dim = 1)
X[:10], y[:10]

train_split = int(0.8 * len(X))
X_train, y_train = X[:train_split], y[: train_split]
X_test, y_test = X[train_split:], y[train_split: ]
len(X_train),len(y_train),len(X_test),len(y_test)

plot_prediction(X_train, y_train, X_test, y_test)
class LinearRegressionModelV2(nn.Module):
  def __init__(self):
    super().__init__()
    #use nn.Linear() for creating the model parameters/also called linear transform/probing layer/fully connected layer
    self.linear_layer = nn.Linear(in_features = 1, out_features = 1)

  def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
    return self.linear_layer(x)

#set the manual seed

torch.manual_seed(42) #the values arent exactly the same because of how pytorch works behind the scenes but, there is still reproducibility
model_1 = LinearRegressionModelV2()
model_1, model_1.state_dict()
model_1.to(device)
next(model_1.parameters()).device
#Set up loss function

loss_fn = nn.L1Loss() #MAE

#Set up optimizer

optimizer = torch.optim.SGD(params = model_1.parameters(),lr = 0.01)

torch.manual_seed(42)
epochs = 200

#put data on the same device as the model

X_train = X_train.to(device)
y_train = y_train.to(device)
X_test = X_test.to(device)
y_test = y_test.to(device)

for epoch in range(epochs):

  model_1.train()

  y_pred = model_1(X_train)

  loss = loss_fn(y_pred, y_train)

  optimizer.zero_grad()

  loss.backward()

  optimizer.step()

  ###Testing
  model_1.eval()

  with torch.inference_mode():
    test_pred = model_1(X_test)
    test_loss = loss_fn(test_pred, y_test)

  if epoch % 10 == 0:
    print(f"Epoch: {epoch} | Loss:{loss} | Test Loss: {test_loss}")

print(model_1.state_dict())

model_1.eval()
with torch.inference_mode():
  y_preds = model_1(X_test)

y_preds

plot_prediction(predictions=y_preds.cpu()) #matplotlib works with numpy which works on cpu, so we need to convert to cpu
</code></pre>
<p>Above is the code snippet! Sorry about the comments, I am still learning Pytorch and these are like my notes</p>

ptrblck: <p><code>plot_prediction</code> expects 4 inputs, which are set to default tensors, and one optional input argument, which is the only passed input argument in:</p>
<pre data-code-wrap="python"><code class="lang-python">plot_prediction(predictions=y_preds.cpu())
</code></pre>
<p>The default arguments are set as <code>X_train</code>, <code>y_train</code>, etc. which are all <code>CUDATensor</code>s, so you would also need to move these to the CPU.<br>
Besides that you are also using <code>X_test</code> directly:</p>
<pre data-code-wrap="python"><code class="lang-python">if predictions is not None:
    plt.scatter(X_test, predictions, c="r", label = "predictions")
</code></pre>
<p>instead of the assigned local variable <code>test_data</code>, which is also wrong.</p>
<p>After fixing this issue and passing all inputs explicitly, your code works.</p>