Hi,
Could you please help me find the mismatch between the CustomDataset and the model?
Thanks!
Code below:
from torch.utils.data import Dataset, DataLoader

import torch.nn as nn

class CustomDataset(Dataset):

    def __init__(self, path,start_idx, end_idx):

        

        self.data = np.load(path)

        self.data = self.data[start_idx:end_idx]

    def __len__(self):

       

        return len(self.data)

    

    def add_noise(self,y):

        newimg = random_shapes((256, 256),min_shapes=30,max_shapes=42,

                       multichannel=False, min_size=20,max_size=30,allow_overlap=True)[0]/255.0

        x = y.copy()

        x[np.where(newimg < 0.9)] = x[np.where(newimg < 0.9)]+1-newimg[newimg < 0.9]

        return x

    def __getitem__(self, idx):

        

        y = self.data[idx]

        x = torch.FloatTensor(self.add_noise(y)).unsqueeze(0)

        

        return x, torch.FloatTensor(y).unsqueeze(0)



class Denoise(nn.Module):

    def __init__(self):

        super(Denoise, self).__init__()

        # encoder layers

        self.enc1 = nn.Conv2d(256, 256, kernel_size=3, padding=1)

        self.enc2 = nn.Conv2d(256, 32, kernel_size=3, padding=1)

        self.enc3 = nn.Conv2d(32, 16, kernel_size=3, padding=1)

        self.enc4 = nn.Conv2d(16, 8, kernel_size=3, padding=1)

        self.pool = nn.MaxPool2d(2, 2)

        

        # decoder layers

        self.dec1 = nn.ConvTranspose2d(8, 8, kernel_size=3, stride=1)  

        self.dec2 = nn.ConvTranspose2d(8, 16, kernel_size=3, stride=1)

        self.dec3 = nn.ConvTranspose2d(16, 32, kernel_size=3, stride=1)

        self.dec4 = nn.ConvTranspose2d(32, 256, kernel_size=3, stride=1)

        self.out = nn.Conv2d(256, 256, kernel_size=3, padding=1)

  

    def forward(self, x):

        # encode

        x = F.relu(self.enc1(x))

        x = self.pool(x)

        x = F.relu(self.enc2(x))

        x = self.pool(x)

        x = F.relu(self.enc3(x))

        x = self.pool(x)

        x = F.relu(self.enc4(x))

        x = self.pool(x) # the latent space representation

        

        # decode

        x = F.relu(self.dec1(x))

        x = F.relu(self.dec2(x))

        x = F.relu(self.dec3(x))

        x = F.relu(self.dec4(x))

        x = F.sigmoid(self.out(x))

 

        return x