Sorry, I missed the use of ParameterDict. The problem then is that you’re doing hstack once, and not in forward(), so it is not in the graph (zero_grad() before the first iteration resets logit_parameter.grad_fn, and it is never re-linked to hstack after that).
So, you’ll have to either do the stacking in forward() or use a stacked parameter, without dictionaries.