Video "PyTorch Fundamentals"

[sakshisaraswat19]

Hello I'm having trouble with training and testing the CNN model.
In 7.2 Training and testing of model_2 using training and testing function:

I've been following the same architecture ,hyperparameters and dataset as in the video, but I noticed my train accuracy is way lower than what's shown. The test accuracy is similar, so I'm wondering what could be causing the difference.

code:

train_step(model= model_2,
             data_loader = train_dataloaders,
             loss_fn = loss_fn,
             optimizer= optimizer,
             accuracy_fn = accuracy_fn,
             device = device)

As for the train_step() function:

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)
  model.train()
   ## Add a loop to loop through the training batches
  for batch, (X,y) in enumerate(data_loader):

    X,y = X.to(device), y.to(device)
    # 1. Forward pass
    y_pred = model(X)

    # 2. Calculate loss and acc(per batch)
    loss = loss_fn(y_pred, y)
    train_loss += loss
    train_acc = accuracy_fn(y_true= y,
                            y_pred = y_pred.argmax(dim=1))

    # 3. optimizer zero grad
    optimizer.zero_grad()

    # 4. loss backward
    loss.backward()

    # 5. optimizer step
    optimizer.step()

   # divide total train loss and accuracy by len of train dataloader
  train_loss /= len(data_loader)
  train_acc /= len(data_loader)
  print(f"Train loss:{train_loss:.5f}| Train accuracy:{train_acc:.2f}%")