well this is the relevant part of the code:
def train_net(net,
              device,
              epochs=5,
              batch_size=2,
              lr=0.001,
              val_percent=0.1,
              save_cp=True,
              img_scale=0.25):

    dataset = BasicDataset(dir_img, dir_mask, img_scale)
    n_val = int(len(dataset) * val_percent)
    n_train = len(dataset) - n_val
    train, val = random_split(dataset, [n_train, n_val])
    train_loader = DataLoader(train, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True)
    val_loader = DataLoader(val, batch_size=batch_size, shuffle=False, num_workers=8, pin_memory=True, drop_last=True)

    writer = SummaryWriter(comment=f'LR_{lr}_BS_{batch_size}_SCALE_{img_scale}')
    global_step = 0

    logging.info(f'''Starting training:
        Epochs:          {epochs}
        Batch size:      {batch_size}
        Learning rate:   {lr}
        Training size:   {n_train}
        Validation size: {n_val}
        Checkpoints:     {save_cp}
        Device:          {device.type}
        Images scaling:  {img_scale}
    ''')

    optimizer = optim.RMSprop(net.parameters(), lr=lr, weight_decay=1e-8, momentum=0.1)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min' , patience=2) # original: 'min' if net.n_classes > 1 else 'max'
    # if net.n_classes > 1:
    #     criterion = nn.CrossEntropyLoss()
    # else:
    #     criterion = nn.CrossEntropyLoss()
    #     # criterion = nn.BCEWithLogitsLoss()
    
    L1_criterion = nn.L1Loss()    
    semantic_criterion = nn.CrossEntropyLoss(reduction='none')
    
    for epoch in range(epochs):
        net.train()
        # epoch_semantic_loss = 0
        epoch_loss = 0
        with tqdm(total=n_train, desc=f'Epoch {epoch + 1}/{epochs}', unit='img') as pbar:
            for batch in train_loader:
                # print('Test Test Test')
                mixedImgs = batch['mixedImage']
                if mixedImgs.size()[0]==1:
                    continue                   
                BackgroundImage = batch['BackgroundImage']
                ReflectionImage = batch['ReflectionImage']
                BackgtoundSemantics = batch['BackgtoundSemantics']
                assert mixedImgs.shape[1] == 3, \
                    f'Network has been defined with {net.n_channels} input channels, ' \
                    f'but loaded images have {mixedImgs.shape[1]} channels. Please check that ' \
                    'the images are loaded correctly.'

                mixedImgs = mixedImgs.to(device=device) #, dtype=torch.float32)
                BackgroundImage = BackgroundImage.to(device=device) #, dtype=torch.float32)*255
                ReflectionImage = ReflectionImage.to(device=device) #, dtype=torch.float32)
                BackgtoundSemantics =  BackgtoundSemantics.to(device=device , dtype=torch.long) # , dtype=torch.long
            
                semantic_pred, BG_img_pred, R_img_pred = net(mixedImgs)
            
                varR = torch.var(torch.abs(R_img_pred-ReflectionImage)) #dim=0, keepdim=True,
                semLos = semantic_criterion(torch.abs(semantic_pred[:,:,:,:]),BackgtoundSemantics[:,0,:,:] )
                varS = torch.var(semLos)
   

the model and other supporting functions are defined, if needed I can post them too.