@ptrblck, I am posting my code for training here…
epochs = 10
steps = 0
running_loss = 0
print_every = 125
train_losses, test_losses = [], []
criterion = nn.MSELoss(reduction='sum')
optimizer = t.optim.Adam(model.parameters(), 5e-4, (0.9, 0.999), eps=1e-08, weight_decay=1e-4)
for epoch in range(epochs):
    for images_train, masks_train in train_loader:
        steps += 1
        images_train = images_train.to(device)
        masks_train = masks_train.type(t.LongTensor)
        #masks_train.unsqueeze_(0)
        #masks_train = masks_train.repeat(3, 1, 1, 1)
        masks_train = masks_train.reshape(masks_train.shape[0], masks_train.shape[2], masks_train.shape[3])
        masks_train = masks_train.to(device)
        optimizer.zero_grad()
        _,logps = model(images_train)
        loss = criterion(logps.float(), masks_train.float())
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        
        if steps % print_every == 0:
            test_loss = 0
            accuracy = 0
            model.eval()
            with t.no_grad():
                for inputs_val, masks_val in test_loader:
                    inputs_val = inputs_val.to(device)
                    masks_val = masks_val.type(t.LongTensor)
                    #masks_val.unsqueeze_(0)
                    #masks_val = masks_val.repeat(3, 1, 1, 1)
                    masks_val = masks_val.reshape(masks_val.shape[0], masks_val.shape[2], masks_val.shape[3])
                    masks_val = masks_val.to(device)
                    _,logps1 = model(inputs_val)
                    batch_loss = criterion(logps1.float(), masks_val.float())
                    test_loss += batch_loss.item()
                    
            train_losses.append(running_loss/len(train_loader))
            test_losses.append(test_loss/len(test_loader))                    
            print(f"Epoch {epoch+1}/{epochs}.. "
                  f"Train loss: {running_loss/print_every:.3f}.. "
                  f"Test loss: {test_loss/len(test_loader):.3f}.. ")
            running_loss = 0
            model.train()

I am getting nan values for training and validation losses. Please let me know if you find any problem with this.