@ptrblck
I am trying to apply the Gradient reversal layer (GRL) from  Speaker Invariant Training using GRL and Unsupervised Domain Adaptation by Backpropagation in the SpeakerNet as mentioned below.
The idea is to learn speaker invariant emotion features by adversarial training using GRL i.e to unlearn speaker information.
Is my implementation correct for GRL here and in main training function, do I have to simply sum the losses of EmotionNet and SpeakerNet and then back-propagate ?
class EncoderNet(nn.Module):
    def __init__(self):
        super(Multitask, self).__init__()       
        self.cnn = TDNN()   #Assume Shared Encoder Here 

    def forward(self, input):
       return self.cnn(input)

        
class Multitask(nn.Module):
    def __init__(self):
        super(Multitask, self).__init__()       
        self.enc_net               =  EncoderNet()
        self.emotion_classfier = EmotionClassfier()
        self.speaker_invariant = SpeakerInvariant()

    def forward(self, x):
        input = self.enc_net(x)
        input_emotion = self.emotion_classfier(input)
        input_speaker = self.speaker_invariant(input) 
       """First confusion is here, how shall handle it here and then how to implemenet this loss function  """ 
        return [input_emotion, input_speaker]

class EmotionClassfier(nn.Module):
    def __init__(self):
        super(EmotionClassfier, self).__init__() 
        self.fc1 = nn.Linear(128, 64)
        self.fc2 = nn.Linear(64, 5)     #5 here is number of emotion classes 
        self.relu = nn.ReLU()

    def forward(self, input):
        input = input.view(input.size(0), -1)
        input = self.fc1(input)
        input = self.relu(input)
        input = self.fc2(input)
        return input

class SpeakerInvariant(nn.Module):
    def __init__(self):
        super(SpeakerInvariant, self).__init__()
        self.fc1 = nn.Linear(128, 64)
        self.fc2 = nn.Linear(64, 10)   #10 here is number of speakers/labels
        self.relu = nn.ReLU()
   
    def forward(self, input):
        input_reverse_gradients = ReverseGradients.apply(input, 0.5)  #alpha is 0.5
         input = self.relu(self.fc1(input_reverse_gradients))
         return self.fc2(input)
"""For GRadient Reversal Layer """
from torch.autograd import Function

class ReverseGradients(Function):
    @staticmethod
    def forward(ctx, x, lamda):
        ctx.lamda = lamda
        return x.view_as(x)

    @staticmethod
    def backward(ctx, grad_output):
        output = grad_output.neg() * ctx.lamda
        return output, None