class NetBlock(nn.Module):
    def __init__(self, in_channels=88, out_channels=64, kernel_size=3, padding=1):
        super(NetBlock, self).__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding),
            nn.Conv2d(out_channels, out_channels,kernel_size, padding=padding),
            nn.BatchNorm2d(out_channels),
            nn.LeakyReLU(negative_slope=0.2, inplace=True)
        )

    def forward(self, x):
        return self.layer(x)


class my_model(nn.Module):

    def __init__(self):
        super(my_model, self).__init__()
                
        self.alpha = nn.Parameter(torch.rand(1))
        self.beta = nn.Parameter(torch.rand(1))
        
        self.layer = nn.ModuleList()
        self.pred = nn.ModuleList()

        self.layer.append(NetBlock(2, 64, 1, 0))
        self.pred.append(Pred(64, 1))

        self.layer.append(NetBlock(81, 64, 1, 0))
        self.pred.append(Pred())

        self.layer.append(NetBlock(kernel_size=1, padding=0))
        self.pred.append(Pred())

        # layer3
        self.layer.append(NetBlock())
        self.pred.append(Pred())

        self.layer.append(NetBlock())
        self.pred.append(Pred())

        self.layer.append(NetBlock())
        self.pred.append(Pred())

        self.layer.append(NetBlock())
        self.pred.append(Pred())

        self.layer.append(NetBlock())
        self.pred.append(Pred())

        self.layer.append(NetBlock())
        self.pred.append(Pred())

        self.layer.append(NetBlock())
        self.pred.append(Pred())

        # layer10
        self.layer.append(NetBlock())
        self.pred.append(Pred())

    def forward(self, x):
        feature_map = []
        pred_map = []
        output = []

        feature_map.append(self.layer[0](normalize(x[0])))
        pred_map.append(self.pred[0](feature_map[0]))
        amp = self.alpha*x[0][:, 0, :, :]+(1-self.alpha)*x[0][:, 1, :, :]
        output.append(torch.unsqueeze(amp,1))

        for i in range(1, len(x)):
            img_shape = (x[i].shape[2], x[i].shape[3])
            feature_map.append(self.layer[i](torch.cat([normalize(x[i]),
                                            F.interpolate(feature_map[i-1], img_shape, mode='bilinear'),
                                            F.interpolate(pred_map[i-1], img_shape, mode='bilinear')], 1)))
            pred_map.append(self.pred[i](feature_map[i]))
            amp = self.beta*x[i][:, 0:4, :, :] + (1-self.beta)*x[i][:, 8:12, :, :]
            phase = pred_map[i][:,4:8,:,:]
            output.append(torch.cat([amp, phase],1))
        return output


Though I am able to train a single instance of this model correctly