Background here. I want to implement Neural Style Transfer using Pytorch from scratch (for educational purpose). There are 2 models, which are transformer net (T) and loss net (L). With given input x, I then will compute o = T(x). Then compute loss from L(x, o). But the loss from loss net did not propagate back to transformer net.
I don’t think I understand how pytorch autograd well enough to know what is going wrong. But the code below is what I came up with, and if I set lossNet on requires_grad to False. The code give me a error of loss does not have grad_fn.
Please advised, Thank you
Code
Here are the simplify code of my current situation where Model A is transformer net and Model B is loss net
class ModelA(nn.Module):

    def __init__(self, requires_grad=True):
        super(ModelA, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 16, 3, padding=1)
        self.relu1 = torch.nn.ReLU(inplace=False)
        self.maxpool = torch.nn.MaxPool2d(2, padding=0)
        self.conv2 = torch.nn.Conv2d(16, 16, 1)
        self.relu2 = torch.nn.ReLU(inplace=False)
        self.conv3T = torch.nn.ConvTranspose2d(16, 3, 1)
        self.relu3 = torch.nn.ReLU(inplace=False)
        self.upSampling = torch.nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.sigmoid = torch.nn.Sigmoid()

    def forward(self, in_x):
        x = in_x
        x = self.relu1(self.conv1(x))
        x = self.maxpool(x)
        x = self.relu2(self.conv2(x))
        x = self.relu3(self.conv3T(x))
        x = self.sigmoid(self.upSampling(x))
        return x

class ModelB(nn.Module):

    def __init__(self):
        super(ModelB, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 64, 3, padding=1)
        self.relu1 = torch.nn.ReLU(inplace=False)
        self.maxpool1 = torch.nn.MaxPool2d(2, padding=1)
        self.conv2 = torch.nn.Conv2d(64, 128, 3, padding=1)
        self.relu2 = torch.nn.ReLU(inplace=False)
        self.maxpool2 = torch.nn.MaxPool2d(2, padding=1)
        self.conv3 = torch.nn.Conv2d(128, 256, 3, padding=1)
        self.relu3 = torch.nn.ReLU(inplace=False)
        self.maxpool3 = torch.nn.MaxPool2d(2, padding=1)
    
    def forward(self, in_x, in_o):
        x = in_x
        o = in_o

        x = self.maxpool1(self.relu1(self.conv1(x)))
        x = self.maxpool2(self.relu2(self.conv2(x)))
        x = self.maxpool3(self.relu3(self.conv3(x)))

        o = self.maxpool1(self.relu1(self.conv1(o)))
        o = self.maxpool2(self.relu2(self.conv2(o)))
        o = self.maxpool3(self.relu3(self.conv3(o)))
        
        return x, o

The Model B is then separate into subnetwork and wrapped inside another class before merge again
class ModuleWrapper(nn.Module):

    def __init__(self, subnetwork, isUseCuda):

        super(ModuleWrapper, self).__init__()

        self.layers = list(subnetwork.children())
        self.net = nn.Sequential(*self.layers)

        self.loss = torch.tensor(0.0).float()
        self.not_inplace = lambda layer: nn.ReLU(inplace=False) if isinstance(layer, nn.ReLU) else layer

        if isUseCuda:
            self.loss = self.loss.cuda()
        
    def forward(self, in_x, in_o):
        if torch.cuda.is_available():
            x = in_x.cuda()
            o = in_o.cuda()
        else:
            x = in_x
            o = in_o
        
        for layer in self.layers:
            layer = self.not_inplace(layer)
            x = layer(x)
            o = layer(o)
            if torch.cuda.is_available():
                x = x.cuda()
                o = o.cuda()

        self.loss = F.mse_loss(x, o)
        return x, o

class MergeWrapper(nn.Module):
    def __init__(self, modules):
        super(MergeWrapper, self).__init__()
        self.net = nn.Sequential(*modules)
    
    def forward(self, in_x, in_o):
        x, o = in_x, in_o
        for module in self.net:
            x, o = module(x, o)
        return x, o
    
    def get_module(self, index):
        assert 0 <= index < len(self.net)
        return self.net[index]
    
    def max_seq(self):
        return len(self.net)

backboneNet= ModelA()
for param in backboneNet.parameters():
    param.requires_grad = True

lossNet = ModelB()
# Assume Model B is pretrained
for param in lossNet.parameters():
    param.requires_grad = False

layers = list(lossNet.children())
subnet = []
subnet.append(ModuleWrapper(nn.Sequential(*layers[0:3]), useCuda))
subnet.append(ModuleWrapper(nn.Sequential(*layers[3:6]), useCuda))
subnet.append(ModuleWrapper(nn.Sequential(*layers[6:9]), useCuda))

lossNet = MergeWrapper(subnet)

Then I test with the following code
testInput = torch.rand(1, 3, 8, 8)

def train(testInput, num_steps = 10):
    optimizer = optim.LBFGS([testInput.requires_grad_()])
    run = [0]
    
    while run[0] <= num_steps:
        def closure():
            optimizer.zero_grad()
            x = testInput
            o = backboneNet(x)
            
            x = x.data.clamp_(0, 1)
            o = o.data.clamp_(0, 1)
            
            lossNet(x, o)
            
            loss = torch.tensor(0.0).float().to(device)
            
            for module in subnet:
                loss += module.loss
            print(loss)
            print_backprop(loss)
            
            loss.backward(retain_graph=True)
            run[0] += 1
            return loss
        
        optimizer.step(closure)

train(testInput)

The printed loss value is not zero when set requires_grad = True on lossNet but the printed back-propagated layers is incorrect. And I have no idea how to fix the issue.