Solved by Balamurali_M in post #2 
              
              
                You can use dataloader. Set the dataloader in such a way that it can return image, binary mask and image filename. I am assuming there are two folders, one for image and other for masks and both have same filenames. 
class DataLoaderEx(Dataset):
    
    def __init__(self,img_path,mask_path):
      …