Solved by duyducvo4444 in post #10 
              
              
                If you are trying to make x to be the same as target ? 
Why don’t you just 
def forward(self, features,targets):
        x = self.model.head(features)
        x = self.last(x)
        x = self.depth1(x)
        x = self.depth2(x)
        loss = some_loss_function(x, targets)
        return x, loss

T…