I have been trying to run this code, but I am having always the same problem.

image1276×549 37.5 KB

disc_I = Discriminator(in_channel=3).to(cfg.DEVICE1)
disc_L = Discriminator(in_channel=3).to(cfg.DEVICE1)
gen_L = Generator(im_channel=3, num_residuals=9).to(cfg.DEVICE1)
gen_I = Generator(im_channel=3, num_residuals=9).to(cfg.DEVICE1)
opt_disc = torch.optim.Adam(
    list(disc_I.parameters()) + list(disc_L.parameters()),
    lr=cfg.LEARNING_RATE,
    betas=(0.5, 0.999),
)
opt_gen = torch.optim.Adam(
    list(gen_I.parameters()) + list(gen_L.parameters()),
    lr=cfg.LEARNING_RATE,
    betas=(0.5, 0.999),
)

L1 = nn.L1Loss()
mse = nn.MSELoss()
g_scaler = torch.cuda.amp.GradScaler()

        # train generator
        with torch.cuda.amp.autocast():
            # adversarial loss for both generators
            D_I_fake = disc_I(fake_initial)
            D_L_fake = disc_L(fake_last)
            G_I_loss = mse(D_I_fake, torch.ones_like(D_I_fake))
            G_L_loss = mse(D_L_fake, torch.ones_like(D_L_fake))
            
            
            #cycle loss
            cycle_last = gen_L(fake_initial)
            cycle_initial = gen_I(fake_last)
            cycle_last_loss = l1(last, cycle_last)
            cycle_initial_loss = l1(initial, cycle_initial)
            
            
            
            #identity loss
            identity_last = gen_L(last)
            identity_initial = gen_I(initial)
            
            identity_last_loss = l1(last, identity_last)
            identity_initial_loss = l1(initial, identity_initial)
            
            
            
            # add all toghether
            G_loss = (G_L_loss + G_I_loss 
                      + cycle_last_loss*cfg.LAMBDA_CYCLE
                      + cycle_initial_loss*cfg.LAMBDA_CYCLE 
                      + identity_last_loss*cfg.LAMBDA_IDENTITY
                      + identity_initial_loss*cfg.LAMBDA_IDENTITY)
        
        opt_gen.zero_grad()
        g_scaler.scale(G_loss)
        g_scaler.step(opt_gen)
        #opt_gen.step()
        g_scaler.update()


Where Generator and discriminator are as following.
class Generator(nn.Module):
    def __init__(self, im_channel, nfg = 64, num_residuals=9):
        super().__init__()
        self.initial = nn.Sequential(
            nn.Conv2d(im_channel, nfg, kernel_size=7, stride=1, padding=3),
            nn.ReLU(inplace=True),
            )
        self.down_block = nn.ModuleList([
            ConvBlock(nfg, nfg*2, kernel_size=3, stride=2, padding=1),
            ConvBlock(nfg*2, nfg*4, kernel_size=3, stride=2, padding=1),
        ])
        
        self.residual_block = nn.Sequential(
            *[ResidualBlock(nfg*4) for _ in range(num_residuals)]
        )
        
        self.up_block = nn.ModuleList([
            ConvBlock(nfg*4, nfg*2, down=False, kernel_size=3, stride=2, padding=1, output_padding=1),
            ConvBlock(nfg*2, nfg, down=False, kernel_size=3, stride=2, padding=1, output_padding=1),
        ])
        
        self.last = nn.Conv2d(nfg, im_channel, kernel_size=7, stride=1, padding=3)
        
    def forward(self, x):
        x = self.initial(x)
        for layer in self.down_block:
            x = layer(x)
        x = self.residual_block(x)
        for layer in self.up_block:
            x = layer(x)
        return torch.tanh(self.last(x))

class Discriminator(nn.Module):
    def __init__(self, in_channel=3, ngf = 64):
        super().__init__()
        
        net = []
        net.append(sub_dis(in_channel, ngf, init=True))
        
        ini = ngf
        n_layer = 3
        for i in range(n_layer):
            end = ngf*2
            net.append(sub_dis(ini, end, stride=1 if i == n_layer-1 else 2))
            ini = end
        
        net.append(nn.Conv2d(ini, 1, kernel_size=4, stride=1, padding=1))
        #net.append(nn.Sigmoid())
        self.model = nn.Sequential(*net)
        
    def forward(self, x):
        return self.model(x)

I could not find any answer on the previous topics that could solve my problem.
Thanks!