import time

import torch
import torch.nn as nn

import numpy as np


class A(nn.Module):
    def __init__(self):
        super(A, self).__init__()
        self.last = nn.Sequential(
            nn.Linear(11, 2),
            nn.LeakyReLU(),
            nn.Linear(2, 2),
            nn.LeakyReLU(),
            nn.Linear(2, 1)
        )

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


class B(nn.Module):
    def __init__(self):
        super(B, self).__init__()
        self.last = nn.Sequential(
            nn.Linear(11, 2),
            nn.LeakyReLU(),
            nn.Linear(2, 2),
            nn.LeakyReLU(),
            nn.Linear(2, 1),    
        )

    def forward(self, x):
        x = self.last(x)
        return x


class TestModule():
    def __init__(self):
        self.a = A()
        self.b = B()

        #self.optimizer = torch.optim.Adam([{"params":self.a.parameters()}, {"params":self.b.parameters()}])
        self.a_optimizer = torch.optim.Adam(self.a.parameters())
        self.b_optimizer = torch.optim.Adam(self.b.parameters())


tm = TestModule()

for _ in range(10000):
    input = np.array([1 for _ in range(11)])

    input_tensor = torch.from_numpy(input).float()

    result_a = tm.a(input_tensor)
    result_b = tm.b(input_tensor)

    loss_a = (result_a-1)**2
    loss_b = (result_b-1)**2

    loss = (loss_a - loss_b)**2

    a_loss = -loss
    b_loss = loss

    tm.a_optimizer.zero_grad()  
    tm.b_optimizer.zero_grad()
    
    # for e in tm.a.parameters():
    #     e.requires_grad = True
    # for e in tm.b.parameters():
    #     e.requires_grad = False

    (a_loss+b_loss).backward()
    #a_loss.backward()
    print("a_backward")
    print(f"a_loss : {a_loss}")
    print(tm.a.last[0].weight.grad)
    print(tm.b.last[0].weight.grad)
    
    # for e in tm.a.parameters():
    #     e.requires_grad = False
    # for e in tm.b.parameters():
    #     e.requires_grad = True
    
    #b_loss.backward()
    print("b_backward")
    print(f"b_loss : {b_loss}")
    print(tm.a.last[0].weight.grad)
    print(tm.b.last[0].weight.grad)

    tm.a_optimizer.step()  
    tm.b_optimizer.step()
    
    print(f"a : {result_a}")
    print(f"b : {result_b}")

    time.sleep(5)

(Code for testing purposes)
I tried various things, but backprop doesn’t seem to work.