01_PyTorch Workflow (Runtime Error: Expected all tensors to be on the same device)

[allen-ajith]

In the last portion of section 01 PyTorch workflow fundamentals, I get an error in the training loop.

RuntimeError                              Traceback (most recent call last)
<ipython-input-10-f6573d2656f5> in <cell line: 20>()
     22   model_1.train()
     23 
---> 24   y_pred = model_1(X_train)
     25 
     26   training_loss = lossfn(y_pred, y_train)

5 frames
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/linear.py in forward(self, input)
    114 
    115     def forward(self, input: Tensor) -> Tensor:
--> 116         return F.linear(input, self.weight, self.bias)
    117 
    118     def extra_repr(self) -> str:

**RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!** (when checking argument for argument mat1 in method wrapper_CUDA_addmm)

I have already put both model_1 and X_train to cuda before starting the training loop.

#Creating a model by subclassing nn.Module

class LinearRegressionModelV2(nn.Module):
  def __init__(self):
    super().__init__()
    #Use nn.Linear() for initializing model parameters(aka: linear transforms/dense layer, probing layer, fully connected layer)
    self.linear_layer = nn.Linear(in_features=1,out_features=1) #how many features are going into the model and what the model outputs

  def forward(self, x: torch.Tensor) -> torch.Tensor:
    return self.linear_layer(X)

#Creating instance of model

torch.manual_seed(42)

model_1=LinearRegressionModelV2()
model_1.state_dict()

model_1.to(device)
next(model_1.parameters()).device


#Training loop

torch.manual_seed(42)

#Write device agnostic code (Put the model and tensors into your GPU)

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



epochs=200

epoch_count=[]
train_loss=[]
test_loss=[]

for epoch in range(epochs):

  model_1.train()
  
  y_pred = model_1(X_train)

  training_loss = lossfn(y_pred, y_train) 

  optimizer.zero_grad()

  training_loss.backward()

  optimizer.step()


  model_1.eval()

  with torch.inference_mode():

    test_pred = model_1(X_test)

    testing_loss = lossfn(test_pred, y_test)

    epoch_count.append(epoch)
    train_loss.append(training_loss.detach())
    test_loss.append(testing_loss)

    if epoch%10 == 0:

      print(f"{epoch} Epochs|| Train Loss:{training_loss} || Test Loss:{testing_loss}")

While printing out X_train.device and next(model_1.parameters()).device the output is device(type='cuda', index=0)

[allen-ajith]

Solved!

It was a typo in forward() method while creating model

def forward(self, x: torch.Tensor) -> torch.Tensor: return self.linear_layer(X)

Correction is

def forward(self, x: torch.Tensor) -> torch.Tensor: return self.linear_layer(x)