Aren’t the parameters that construct self.logit_parameter being registered with self.params = torch.nn.ParameterDict(parameters) ?
But if I understand you correctly this is not enough as the module does not know that the registered self.params are the same parameters that constructed self.logit_parameter. Basically, what I got now is that the module knows about the self.logit_parameter and the self.params that constructed self.logit_parameter but it does not know anything about how self.logit_parameter was constructed (the computation graph is not visible for the module).
Maybe, to formulate my initial question more clearly:
Is there a way to register the computation graph that constructs self.logit_parameter without  wrapping the computation graph construction in module of its own?