Thank you for the reply
I had came up with somewhat adhoc idea
class OutputProxy(nn.Module):
    states = []
    def __init__(self):
        super(OutputProxy, self).__init__()
        
    def forward(self, x):
        OutputProxy.push(x.detach())
        return x
        
    @staticmethod
    def push(item):
        OutputProxy.states.append(item)
    
    @staticmethod
    def pop():
        return OutputProxy.states.pop()

If I add OutputProxy layers after the desire output, and then call OutputProxy.pop() in the custom loss function. If I do this, does it have any caveats like this can be done only on cpu or something like that?
And if I could, would you recommend it?