As a followup, here are the model definitions:
class F_extractor(nn.Module):

def __init__(self, d):
    super(F_extractor, self).__init__()
    self.fc1 = nn.Linear(2,d)
    
def forward(self, x):
    x = F.relu(self.fc1(x))
    return x

class Classifier(nn.Module):
def __init__(self, d):
    super(Classifier, self).__init__()
    self.fc1 = nn.Linear(d, 2)
    

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

class GradReverse(Function):
def __init__(self, lambd):
    self.lambd = lambd
    
def forward(self, x):
    return x.view_as(x)

def backward(self, grad_output):
    return (grad_output * (-self.lambd))

def grad_reverse(x, lambd):
return GradReverse(lambd)(x)

class Discriminator(nn.Module):
def __init__(self, d, lambd = 0):
    super(Discriminator, self).__init__()
    self.fc1 = nn.Linear(d, 8) 
    self.fc2 = nn.Linear(8, 2)
    self.lambd = lambd
    
def forward(self, x):
    x = grad_reverse(x, self.lambd)
    x = self.fc1(x)
    x = F.relu(x)
    x = self.fc2(x)

    return x