well, it dosnt work yet.
here is the code fully code (it is pretty amateurish ):
import argparse
import logging
import os
import sys
import cv2 

import SSIM.pytorch_ssim as ssim
import torchvision
# import numpy as np
import torch
import torch.nn as nn
from torch import optim
from tqdm import tqdm
from numpy import linalg as LA

from eval import eval_net
from unet import UNet, DeepLabv3_plus, deeplab_v3_separation
# import coco_annotation

from torch.utils.tensorboard import SummaryWriter
from utils.dataset import BasicDataset
from torch.utils.data import DataLoader, random_split

dir_img = 'data/train2017/train2017_640480/'
dir_mask = 'data/train2017/train2017_640480_anns/'
dir_checkpoint = 'checkpoints/'
# os.environ['CUDA_LAUNCH_BLOCKING'] = "1"

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 = -BackgtoundSemantics*torch.log10(semantic_pred)-(1-BackgtoundSemantics)*torch.log10(1-semantic_pred)
                semLos = semantic_criterion(torch.abs(semantic_pred[:,:,:,:]),BackgtoundSemantics[:,0,:,:] )
                varS = torch.var(semLos)
                
                loss = L1_criterion(BG_img_pred, BackgroundImage)
                # loss += 0.0003*LA.norm(cv2.Canny(BG_img_pred,100,200) - cv2.Canny(BackgroundImage,100,200))
                loss += 0.6*(1 - ssim.ssim(BG_img_pred, BackgroundImage))
                loss += 0.8*L1_criterion(R_img_pred, ReflectionImage)
                loss = loss/(2*varR.square())
                loss += semantic_criterion(semantic_pred[:,:,:,:],BackgtoundSemantics[:,0,:,:])/(2*varS.square())
                epoch_loss += loss.item()
                
            
                writer.add_scalar('Loss/train', loss.item(), global_step)
                pbar.set_postfix(**{'loss (batch)': loss.item()})

                optimizer.zero_grad()
                loss.backward()
                # semantic_loss.backward()
                nn.utils.clip_grad_value_(net.parameters(), 0.1)
                optimizer.step()

                pbar.update(mixedImgs.shape[0])
                global_step += 1
                if global_step % (len(dataset) // (10 * batch_size)) == 0:
                    for tag, value in net.named_parameters():
                        tag = tag.replace('.', '/')
                        writer.add_histogram('weights/' + tag, value.data.cpu().numpy(), global_step)
                        writer.add_histogram('grads/' + tag, value.grad.data.cpu().numpy(), global_step)
                    val_score = eval_net(net, val_loader, device)
                    # scheduler.step(val_score)
                    writer.add_scalar('learning_rate', optimizer.param_groups[0]['lr'], global_step)

                    if 1 > 1: #net.n_classes
                        logging.info('Validation cross entropy: {}'.format(val_score))
                        writer.add_scalar('Loss/test', val_score, global_step)
                    else:
                        logging.info('Validation Dice Coeff: {}'.format(val_score))
                        writer.add_scalar('Dice/test', val_score, global_step)

                    writer.add_images('images', mixedImgs, global_step)
                    # if net.n_classes == 1:
                    #     writer.add_images('masks/true', BG_img_pred, global_step)
                    #     writer.add_images('masks/pred', torch.sigmoid(imgs_pred) > 0.5, global_step)

        if save_cp:
            try:
                os.mkdir(dir_checkpoint)
                logging.info('Created checkpoint directory')
            except OSError:
                pass
            torch.save(net.state_dict(),
                       dir_checkpoint + f'CP_epoch{epoch + 1}.pth')
            logging.info(f'Checkpoint {epoch + 1} saved !')

    writer.close()


def get_args():
    parser = argparse.ArgumentParser(description='Train the UNet on images and target masks',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('-e', '--epochs', metavar='E', type=int, default=5,
                        help='Number of epochs', dest='epochs')
    parser.add_argument('-b', '--batch-size', metavar='B', type=int, nargs='?', default=2,
                        help='Batch size', dest='batchsize')
    parser.add_argument('-l', '--learning-rate', metavar='LR', type=float, nargs='?', default=0.1,
                        help='Learning rate', dest='lr')
    parser.add_argument('-f', '--load', dest='load', type=str, default=False,
                        help='Load model from a .pth file')
    parser.add_argument('-s', '--scale', dest='scale', type=float, default=0.25,
                        help='Downscaling factor of the images')
    parser.add_argument('-v', '--validation', dest='val', type=float, default=10.0,
                        help='Percent of the data that is used as validation (0-100)')

    return parser.parse_args()


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
    args = get_args()
    # args = parser.parse_args()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    logging.info(f'Using device {device}')

    # Change here to adapt to your data
    # n_channels=3 for RGB images
    # n_classes is the number of probabilities you want to get per pixel
    #   - For 1 class and background, use n_classes=1
    #   - For 2 classes, use n_classes=1
    #   - For N > 2 classes, use n_classes=N
    net = deeplab_v3_separation()  #UNet(n_channels=3, n_classes=3, bilinear=True)
    # logging.info(f'Network:\n'
    #              f'\t{net.n_channels} input channels\n'
    #              f'\t{net.n_classes} output channels (classes)\n'
    #              f'\t{"Bilinear" if net.bilinear else "Transposed conv"} upscaling')

    if args.load:
        net.load_state_dict(
            torch.load(args.load, map_location=device)
        )
        logging.info(f'Model loaded from {args.load}')

    net.to(device=device)
    # faster convolutions, but more memory
    # cudnn.benchmark = True

    try:
        train_net(net=net,
                  epochs=args.epochs,
                  batch_size=args.batchsize,
                  lr=args.lr,
                  device=device,
                  img_scale=args.scale,
                  val_percent=args.val / 100)
    except KeyboardInterrupt:
        torch.save(net.state_dict(), 'INTERRUPTED.pth')
        logging.info('Saved interrupt')
        try:
            sys.exit(0)
        except SystemExit:
            os._exit(0)

this is the model:
# -*- coding: utf-8 -*-
"""
Created on Fri Aug 28 14:09:04 2020

@author: aviel
"""


import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.model_zoo as model_zoo
import torchvision.models as models

class deeplab_v3_separation(nn.Module):    
    def __init__(self):
        super(deeplab_v3_separation, self).__init__()
       


        deepLabV3ResNet101 = models.segmentation.deeplabv3_resnet101(pretrained=True)
        features_convs_and_ASSP = list(deepLabV3ResNet101.children())
        
        ResNet101 = models.resnet101(pretrained=True)
        ResNet101_features = list(ResNet101.children())[:-2]

        ASSP_layers = models.segmentation.deeplabv3_resnet101(pretrained=True)
        ASSP_features = list(ASSP_layers.children())
        
        self.ResNet101_features = nn.Sequential(*ResNet101_features)
        self.ASSP_features = nn.Sequential(*ASSP_features[1])
        self.Semantic_convs = nn.Conv2d(2069, 92, 1)
        self.semanticSoftmax = nn.Softmax2d()
        self.Background_convs = nn.Conv2d(2164, 3, 1)
        self.Separation_convs = nn.Conv2d(2167, 3, 1)
        # self.up = nn.ConvTranspose2d(2069 , 2069, kernel_size=2, stride=2)


        
    def forward(self, x):
        input_shape = x.shape[-2:]
        x1 = self.ResNet101_features(x)
        x2 = self.ASSP_features(x1)
        
        x3 = torch.cat((x1,x2),1)
        x3 = F.interpolate(x3, size=input_shape, mode='bilinear', align_corners=False)

        y1 = self.Semantic_convs(x3)
        # y1 = self.semanticSoftmax(y1)
        y2 = self.Background_convs(torch.cat((x,x3,y1),1))
        y3 = self.Separation_convs(torch.cat((x,y1,y2,x3),1))
        
        return y1, y2, y3

and the train loader:
import os
from os.path import splitext
from os import listdir
import numpy as np
from glob import glob
import torch
from torch.utils.data import Dataset
import logging
from PIL import Image

class BasicDataset(Dataset):
    def __init__(self, imgs_dir, masks_dir, scale=1):
        self.imgs_dir = imgs_dir
        self.masks_dir = masks_dir
        self.scale = scale
        assert 0 < scale <= 1, 'Scale must be between 0 and 1'

        self.ids = [splitext(file)[0] for file in listdir(imgs_dir)
                    if not file.startswith('.')]
        logging.info(f'Creating dataset with {len(self.ids)} examples')

    def __len__(self):
        return len(self.ids)

    @classmethod
    def preprocess(cls, pil_img, scale):
        w, h = pil_img.size
        newW, newH = int(scale * w), int(scale * h)
        assert newW > 0 and newH > 0, 'Scale is too small'
        pil_img = pil_img.resize((newW, newH))

        img_nd = np.array(pil_img)

        if len(img_nd.shape) == 2:
            img_nd = np.expand_dims(img_nd, axis=2)

        # HWC to CHW
        img_trans = img_nd.transpose((2, 0, 1))
        if img_trans.max() > 1:
            img_trans = img_trans / 255

        return img_trans

    def semantic_preprocess(cls, pil_img, scale):
        w, h = pil_img.size
        newW, newH = int(scale * w), int(scale * h)
        assert newW > 0 and newH > 0, 'Scale is too small'
        pil_img = pil_img.resize((newW, newH))

        img_nd = np.array(pil_img)

        if len(img_nd.shape) == 2:
            img_nd = np.expand_dims(img_nd, axis=2)

        # HWC to CHW
        img_trans = img_nd.transpose((2, 0, 1))
        # if img_trans.max() > 1:
        #     img_trans = img_trans / 255

        return img_trans

    def __getitem__(self, i):
        idx = self.ids[i]
        idx2 = self.ids[int(len(self.ids)*torch.rand(1))]
        img_file = glob(self.imgs_dir + idx + '.*')
        img_file2 = glob(self.imgs_dir + idx2 + '.*')
        assert len(img_file) == 1, \
            f'Either no mask or multiple masks found for the ID {idx}: {img_file}'
        assert len(img_file2) == 1, \
            f'Either no image or multiple images found for the ID {idx}: {img_file2}'
        img = Image.open(img_file[0])
        img2 = Image.open(img_file2[0])
        
        while img.mode == 'L' :
            i+=1
            idx = self.ids[i]
            img_file = glob(self.imgs_dir + idx + '.*')
            assert len(img_file) == 1, \
                f'Either no mask or multiple masks found for the ID {idx}: {img_file}'
            img = Image.open(img_file[0])
        
        while img2.mode == 'L' :
            idx2 = self.ids[int(len(self.ids)*torch.rand(1))]
            img_file2 = glob(self.imgs_dir + idx2 + '.*')
            assert len(img_file2) == 1, \
                f'Either no image or multiple images found for the ID {idx}: {img_file2}'
            img2 = Image.open(img_file2[0])
        
        assert img.size == img2.size, \
            f'Image and mask {idx} should be the same size, but are {img.size} and {img2.size}'
        
        
        img = self.preprocess(img, self.scale)
        img2 = self.preprocess(img2, self.scale)

        mask_file = glob(self.masks_dir + idx + '.*')
        BackgtoundSemantics = Image.open(mask_file[0])
        BackgtoundSemantics = self.semantic_preprocess(BackgtoundSemantics, self.scale)
        # BackgtoundSemantics = np.array(BackgtoundSemantics)
        if BackgtoundSemantics.max() > 1:
            BackgtoundSemantics = BackgtoundSemantics-92

        # img = np.array(img)/255
        # img2 = np.array(img2)/255
                
        # outPutImages = (np.append(img,img2,axis=0))
        
        BackgroundCoeff = 0.5+0.1*np.random.randint(5,size=1)
        ReflectionCoeff = 1-BackgroundCoeff
        mixedImage = np.array(BackgroundCoeff*img + ReflectionCoeff*img2)

        return {
            'mixedImage': torch.from_numpy(mixedImage).type(torch.float32),
            'BackgroundImage': torch.from_numpy(img).type(torch.float32),
            'ReflectionImage': torch.from_numpy(img2).type(torch.float32),
            'BackgtoundSemantics': torch.from_numpy(BackgtoundSemantics).type(torch.float32)
        }