Section 5 (PyTorch computer vision): I am getting nan values for train and test loss (lessons 111 and 112)

[AxisMeetsWorld]

😕 Hi all, I tried to complete the testing/training loops for batched data and am getting some odd "nan" values when I am looking at the train loss and test loss that are to accumulate through the code. I think I traced it back to model_0(X) but I can't understand why I would be getting different results than @mrdbourke here where around 3.3 you are getting training data. I was able to trace it down somewhat by printing out y_pred each 200 batches. # if batch %200 ==0: # print(f"y_pred: {y_pred}") When I printed that out, I got a tensor loaded with "nan" values.
I went back to look at the definition of my model_0 where it was last defined and I checked that it was giving the same results you had after the flattening, and that was consistent. model_0 = FashionMNISTModelV0( input_shape = 784, #28*28 hidden_units = 10, #how many units are in the hidden layer output_shape= len(class_names) ).to("cpu")
So at this point, I am thoroughly confused as to why I am getting these "nan" values. 😕
Here is my loop code in case there is anything I am missing:

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

#Set the number of epochs for training time (keep small for training)
epochs = 3

#Create training and test loops
for epoch in tqdm(range(epochs)):
  print(f"Epoch: {epoch}\n----")
  #Training 
  train_loss = 0
  #Add a loop to loop through batches (recall that the dataloader was what separated all into batches)
  for batch, (X, y) in enumerate(train_dataloader): #Note the 2 different parameters passed to this for loop
    model_0.train()
    #1. Forward pass
    y_pred = model_0(X)

    #2. Calculate the loss per batch
    loss = loss_fn(y_pred, y)
    train_loss += loss #accumulate the training loss

    # Here is where I tried to see why I was getting nan values.
    # if batch %200 ==0:
      # print(f"y_pred: {y_pred}")

    #3. Optimize zero grad
    optimizer.zero_grad

    #4. Loss backward
    loss.backward()

    #5 Optimizer step
    optimizer.step()

    # Print things out
    if batch %400 == 0:
      print(f"Looked at: {batch * len(X)/len(train_dataloader.dataset)} samples.")
  #Divide our train loss by the length of the data loader:
  train_loss /= len(train_dataloader)

  #Testing
  test_loss, test_acc = 0,0 
  model_0.eval()
  with torch.inference_mode():
    for X, y in test_dataloader:
      #Forward pass
      test_pred = model_0(X)

      #calculate loss
      test_loss += loss_fn(test_pred, y)

      #calculate the accuracy (preds need to be the same as y_true)
      test_acc += accuracy_fn(y_true=y, y_pred=test_pred.argmax(dim =1))

    #Calculate the test loss ave per batch by dividing test loss by length of test dataloader per batch
    test_loss /= len(test_dataloader)

    #Calculate the test acc per batch
    test_acc /= len(test_dataloader)

  #Print results
  print(f"\nTrain Loss: {train_loss:.4f} |Test loss: {test_loss:.4f} |Test acc: {test_acc:.4f}")

  #calculate the training time
  train_time_end_on_cpu = timer()
  total_train_time_model_0 = print_train_time(start=train_time_start_on_cpu, 
                                              end= train_time_end_on_cpu,
                                              device=str(next(model_0.parameters()).device))```

[AxisMeetsWorld]

😆 I think I figured it out! The first time I ran it, I did run from the beginning but my code indents were off. It is a little weird, but the second time, I had to rerun the cell with the data loader definition. It worked, and then the "X" that was passed into model_0, populated. I don't think it was populating because somehow the dataloader wasn't being recognized. I have it working now. 😏

[mrdbourke]

Nice work @AxisMeetsWorld !

You wouldn't believe how many times that's happened to me when coding in Google Colab...

I'd be staring at a training/testing loop for an hour before the error revealed itself.

Glad to hear you fixed it!