PyTorch for Deep Learning Bootcamp - Device issue

[Packetouille]

I have run into a device error that for the life of me I cannot understand. I am following along with the course but have run into the following device issue. Data is on the "cpu", but we have device agnostic code for the data (and for the model though we send that to the device during during instantiation) within the TRAIN & TEST METHODS.

I tested with print statements that both the model and data are on the gpu during the training step. Yet I am still receiving the error that says "Expected all tensors to be on the same device...." If i'm reading the error correctly, the issue seems to happen during the loss function call (using nn.CrossEntropyLoss()). But within the training method, we are calling loss.backward() after the data has been sent to the gpu....

The code below is code that I literally just copied from the course to see if I had typed something wrong, and still I am receiving the same error. Hoping that something will pop out to the more experienced eyes.

# TRAIN METHOD ########################################
def train_step(model: torch.nn.Module,
               data_loader: torch.utils.data.DataLoader,
               loss_fn: torch.nn.Module,
               optimizer: torch.optim.Optimizer,
               accuracy_fn,
               device: torch.device = device):
    train_loss, train_acc = 0, 0

    model.to(device)

    for batch, (X, y) in enumerate(data_loader):
        # Send data to GPU
        X, y = X.to(device), y.to(device)

        # 1. Forward pass
        y_pred = model(X)

        # 2. Calculate loss
        loss = loss_fn(y_pred, y)
        train_loss += loss
        train_acc += accuracy_fn(y_true=y,
                                 y_pred=y_pred.argmax(dim=1)) # Go from logits -> pred labels

        # 3. Optimizer zero grad
        optimizer.zero_grad()

        # 4. Loss backward
        loss.backward()

        # 5. Optimizer step
        optimizer.step()

    # Calculate loss and accuracy per epoch and print out what's happening
    train_loss /= len(data_loader)
    train_acc /= len(data_loader)
    print(f"Train loss: {train_loss:.5f} | Train accuracy: {train_acc:.2f}%")


# TEST METHOD ########################################
def test_step(data_loader: torch.utils.data.DataLoader,
              model: torch.nn.Module,
              loss_fn: torch.nn.Module,
              accuracy_fn,
              device: torch.device = device):
    test_loss, test_acc = 0, 0
    model.to(device)
    model.eval() # put model in eval mode
    # Turn on inference context manager
    with torch.inference_mode(): 
        for X, y in data_loader:
            # Send data to GPU
            X, y = X.to(device), y.to(device)
            
            # 1. Forward pass
            test_pred = model(X)
            
            # 2. Calculate loss and accuracy
            test_loss += loss_fn(test_pred, y)
            test_acc += accuracy_fn(y_true=y,
                y_pred=test_pred.argmax(dim=1) # Go from logits -> pred labels
            )
        
        # Adjust metrics and print out
        test_loss /= len(data_loader)
        test_acc /= len(data_loader)
        print(f"Test loss: {test_loss:.5f} | Test accuracy: {test_acc:.2f}%\n")


# INSTANTIATING MODEL - I don't see a need to post the model , but it is a CNN model ########################
torch.manual_seed(42)
model_2 = FashionMNISTModelV2(input_shape=1, # same as number of input color channels
                              hidden_units=10,
                              output_shape=len(class_names)).to(device) # same as number of classes


# SETUP LOSS AND OPTIMIZER ###############################################
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(params=model_2.parameters(), 
                             lr=0.1)


#TRAINING & TESTING #################################################
# Import tqdm for a progress bar
from tqdm.auto import tqdm
from timeit import default_timer as timer

# Set the seed and start the timer
torch.manual_seed(42)
torch.cuda.manual_seed(42)
train_time_start_on_gpu = timer()

# Set the number of epochs (we'll keep this small for faster training time)
epochs = 3

# Create training and test loop
for epoch in tqdm(range(epochs)): # For tqdm to work, simply wrap the iterator in tqdm()
  print(f"Epoch: {epoch}\n-----")

  train_step(model=model_2,
             data_loader=train_dataloader,
             loss_fn=loss_fn,
             optimizer=optimizer,
             accuracy_fn=accuracy_fn,
             device=device)
  test_step(model=model_2,
            data_loader=test_dataloader,
            loss_fn=loss_fn,
            accuracy_fn=accuracy_fn,
            device=device)

  # Print out what's happening
  # print(f"\nTrain loss: {train_loss:.4f} | Test loss: {test_loss:.4f} | Test acc: {test_acc:.4f}")

# Calculate training time
train_time_end_on_gpu = timer()
total_train_time_model_2 = print_train_time(start=train_time_start_on_gpu,
                                            end=train_time_end_on_gpu,
                                            device=device)

Error that I'm receiving: #################################################

[LuluW8071]

To me the code looks ok however you can check the code by using printing statements as well

# TRAIN METHOD ########################################
def train_step(model: torch.nn.Module,
               data_loader: torch.utils.data.DataLoader,
               loss_fn: torch.nn.Module,
               optimizer: torch.optim.Optimizer,
               accuracy_fn,
               device: torch.device = device):
    train_loss, train_acc = 0, 0

    model.to(device)
    print(model.to(device))

Do same for X, y
Also sometimes I have found restarting colab session resolves the problem itself.