softmax and argmax in training loop why?

[StvMy]

*Fit the model
torch.manual_seed(42)

*Set number of epochs
epochs = 100

*Put data to target device
X_blob_train, y_blob_train = X_blob_train.to(device), y_blob_train.to(device)
X_blob_test, y_blob_test = X_blob_test.to(device), y_blob_test.to(device)


for epoch in range(epochs):
    Training
    model_4.train()

    # 1. Forward pass
    **y_logits = model_4(X_blob_train) # model outputs raw logits 
    y_pred = torch.softmax(y_logits, dim=1).argmax(dim=1) # go from logits -> prediction probabilities -> prediction labels**
    # print(y_logits)
    # 2. Calculate loss and accuracy
    loss = loss_fn(y_logits, y_blob_train) 
    acc = accuracy_fn(y_true=y_blob_train,
                      y_pred=y_pred)

    # 3. Optimizer zero grad
    optimizer.zero_grad()

    # 4. Loss backwards
    loss.backward()

    # 5. Optimizer step
    optimizer.step()

Above are the code from the course I don't really get it why do we need to add softmax and argmax in the forward pass because from what I read online, crossentropyloss took care of it in the background. I test it just with the logits and the result is similar or just a little better then processing the logits, is processing the logits add something to the model maybe less overfitting.

BELOW ARE MY TRAINING RESULT AND LOOP

EPOCH: 299 | train loss: 0.029004432260990143 | test loss: 0.016272680833935738 | Train accuracy 99.38% | Test Accuracy: 100.00%

for epoch in range(epochs):
  multi_class_model.train()

  * Forward pass
  y_preds = multi_class_model(x_train)

  * Calculate loss and accuracy
  loss = loss_fn(y_preds,y_train.long())

  acc = accuracy_fnV2(y_preds,y_train.long())

  loss.backward()

  optimizer.step()

  optimizer.zero_grad()