Thank you! Can I ask if what values should I expect from my output? my training script goes like this:
BATCH_SIZE = 10

EPOCHS = 1

def train(model):

  model.train()

  for epoch in range(EPOCHS):

      for i in tqdm(range(0, len(img_all), BATCH_SIZE)): 

          batch_img_all = img_all[i:i+BATCH_SIZE].view(-1, 3, 224, 224)

          batch_mask_all = mask_all[i:i+BATCH_SIZE].view(-1, 1, 224, 224)

        

          model.zero_grad()

          outputs = model(batch_img_all)

          batch_mask_all = torch.argmax(batch_mask_all, dim=1)

          loss = loss_function(outputs, batch_mask_all)

          loss.backward()

          optimizer.step()    # Does the update

      print(f"Epoch: {epoch}. Loss: {loss}")

  

  return batch_img_all, batch_mask_all, outputs

IMG, MSK, OUTPUT = train(model)