I am running a U-net model, I am getting a Type Error  when training the model, ``
‘’
'num_epochs = 30

#criterion = nn.BCEWithLogitsLoss()

criterion = nn.MSELoss()

optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)

scaler = torch.cuda.amp.GradScaler()

# Decay LR by a factor of 0.1 every several epochs

sch = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma= 0.5)

train(model, criterion, optimizer, num_epochs, scaler,sch)

train_transform = transforms.Compose([

       transforms.Resize((256, 256)),

       #transforms.RandomResizedCrop(256),

       transforms.ToTensor(),

       #transforms.ColorJitter(),

       transforms.RandomHorizontalFlip(),

       transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])

])

test_transform = transforms.Compose([

       transforms.Resize((256,256)),

       #transforms.CenterCrop(256),

       transforms.ToTensor(),

       transforms.Normalize([0.5], [0.5])

   ])

I am guessing that this part of the code is throwing an error.
class Flare(Dataset):

    def __init__(self, flare_dir, wf_dir, flare_img_, wf_img_, transform = None):

        self.flare_dir = flare_dir

        self.wf_dir = wf_dir

        self.transform = transform

        self.flare_l = flare_img_

        self.wf_l = wf_img_

        

    def __len__(self):

        return len(self.flare_l)

    def __getitem__(self, idx):

        wf_img = None

        f_img = Image.open(os.path.join(self.flare_dir, self.flare_l[idx])).convert("RGB")

        for i in self.wf_l:

            if (self.flare_l[idx].split('.')[0][4:] == i.split('.')[0]):

                wf_img = Image.open(os.path.join(self.wf_dir, i)).convert("RGB")

                break

        

        f_img = self.transform(f_img)

        wf_img = self.transform(wf_img)

        

        return f_img, wf_img

flare_dir = '/content/gdrive/MyDrive/data/Flare_img'

wf_dir = '/content/gdrive/MyDrive/data/Without_Flare_'

flare_img = os.listdir(flare_dir)

wf_img = os.listdir(wf_dir)

wf_img.sort()

flare_img.sort()

print(wf_img[0])

train_ds = Flare(flare_dir, wf_dir, flare_img[0:5000], wf_img[0:5000], train_transform)

train_loader = torch.utils.data.DataLoader(dataset=train_ds,

                                           batch_size=BATCH_SIZE, 

                                           shuffle=True)

print(len(train_ds))

This is the part of the code where i am suspecting the error is from, not sure how to tackle it.