Looking at the following code, what changes would you suggest to improve the accuracy of the results: import torch
import torch.nn as nn
import torch.nn.functional as F

class ChannelAttention(nn.Module):
    def __init__(self, in_channels, reduction=16):
        super(ChannelAttention, self).__init__()
        hidden = max(1, in_channels // reduction)
        self.mlp = nn.Sequential(
            nn.Linear(in_channels, hidden, bias=False),
            nn.LeakyReLU(inplace=True),
            nn.Linear(hidden, in_channels, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.shape
        avg = F.adaptive_avg_pool2d(x, 1).view(b, c)
        weights = self.mlp(avg).view(b, c, 1, 1)
        return x * weights

class AxialDepthwiseConv(nn.Module):
    def __init__(self, channels):
        super(AxialDepthwiseConv, self).__init__()
        self.h = nn.Conv2d(channels, channels, kernel_size=(1, 7), padding=(0, 3), groups=channels, bias=False)
        self.v = nn.Conv2d(channels, channels, kernel_size=(7, 1), padding=(3, 0), groups=channels, bias=False)
        self.pw = nn.Conv2d(channels, channels, kernel_size=1, bias=False)
        self.ins = nn.InstanceNorm2d(channels)
        self.relu = nn.LeakyReLU(inplace=True)

    def forward(self, x):
        h = self.h(x)
        v = self.v(x)
        combined = h + v + x
        out = self.pw(combined)
        return self.relu(self.ins(out))

class EncoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(EncoderBlock, self).__init__()
        self.axial_dw = AxialDepthwiseConv(in_channels)
        self.ca = ChannelAttention(in_channels)
        self.pw_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
        self.ins = nn.InstanceNorm2d(out_channels)
        self.relu = nn.LeakyReLU(inplace=True)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

    def forward(self, x):
        x = self.axial_dw(x)
        x = self.ca(x)
        x = self.pw_conv(x)
        x = self.relu(self.ins(x))
        skip = x
        x = self.pool(x)
        return x, skip

class DecoderBlock(nn.Module):
    def __init__(self, in_channels, skip_channels, out_channels):
        super(DecoderBlock, self).__init__()
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.pw_conv1 = nn.Conv2d(in_channels + skip_channels, out_channels, kernel_size=1, bias=False)
        self.axial_dw = AxialDepthwiseConv(out_channels)
        self.ca = ChannelAttention(out_channels)
        self.pw_conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=1, bias=False)
        self.ins = nn.InstanceNorm2d(out_channels)
        self.relu = nn.LeakyReLU(inplace=True)

    def forward(self, x, skip):
        x = self.upsample(x)
        x = torch.cat([x, skip], dim=1)
        x = self.pw_conv1(x)
        x = self.relu(self.ins(x))
        x = self.axial_dw(x)
        x = self.ca(x)
        x = self.pw_conv2(x)
        return self.relu(self.ins(x))

class Bottleneck(nn.Module):
    def __init__(self, channels):
        super(Bottleneck, self).__init__()
        self.pw1 = nn.Conv2d(channels, channels, kernel_size=1, bias=False)
        self.axial_dw1 = AxialDepthwiseConv(channels)
        self.ca1 = ChannelAttention(channels)
        self.axial_dw2 = AxialDepthwiseConv(channels)
        self.ca2 = ChannelAttention(channels)
        self.pw2 = nn.Conv2d(channels, channels, kernel_size=1, bias=False)
        self.ins = nn.InstanceNorm2d(channels)
        self.relu = nn.LeakyReLU(inplace=True)

    def forward(self, x):
        x = self.pw1(x)
        x = self.relu(self.ins(x))
        residual = x
        y = self.axial_dw1(x)
        y = self.ca1(y)
        z = self.axial_dw2(x)
        z = self.ca2(z)
        a = y + z + residual
        a = self.pw2(a)
        a = self.relu(self.ins(a))
        return a


class Generator(nn.Module):
    def __init__(self, in_channels=3, base_channels=16, out_channels=3):
        super(Generator, self).__init__()

        self.InitialConv=nn.Conv2d(in_channels,base_channels,kernel_size=3,padding=1)
        self.enc1 = EncoderBlock(base_channels, base_channels)
        self.enc2 = EncoderBlock(base_channels, base_channels)

        self.bottleneck = Bottleneck(base_channels)

        self.dec1 = DecoderBlock(base_channels, base_channels, base_channels)
        self.dec2 = DecoderBlock(base_channels, base_channels, base_channels)

        self.final_conv = nn.Conv2d(base_channels, out_channels, kernel_size=3, padding=1)

    def forward(self, x):
        x = self.InitialConv(x)
        x, skip1 = self.enc1(x)
        x, skip2 = self.enc2(x)

        x = self.bottleneck(x)

        x = self.dec1(x, skip2)
        x = self.dec2(x, skip1)

        out = self.final_conv(x)
        out = torch.tanh(out)
        return out


class Discriminator(nn.Module):
    def __init__(self, in_channels=3):
        super(Discriminator, self).__init__()

        def down_layer(in_filters, out_filters, normalization=True):
            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]
            if normalization:
                layers.append(nn.InstanceNorm2d(out_filters))
            layers.append(nn.LeakyReLU(0.2, inplace=True))
            return layers

        self.model = nn.Sequential(
            *down_layer(2*in_channels, 64),
            *down_layer(64, 128),
            *down_layer(128, 256),
            nn.ZeroPad2d((1, 1, 1, 1)),
            nn.Conv2d(256, 512, 4, padding=0, bias=False),
            nn.InstanceNorm2d(512),
            nn.ZeroPad2d((1, 1, 1, 1)),
            nn.Conv2d(512, 1, 4, padding=0, bias=False)
        )

    def forward(self, imgA, imgB):
        img_input = torch.cat((imgA, imgB), 1)
        return self.model(img_input)


dummy=torch.rand((1,3,256,256))
dummy1=torch.rand((1,3,256,256))
g_model=Generator()
print(g_model(dummy).shape)
d_model=Discriminator()
print(d_model(dummy,dummy1).shape)