Your code should work and it seems the gradients are just “killed” by the ReLUs in pre_choices.
I.e. if you remove both nn.ReLU() non-linearities, you’ll see valid updated, as they might create zero gradients for some iterations.
This doesn’t necessarily mean that your model won’t get any valid gradients in the following iterations.
Also, if you want to use an integer index instead of keys, you could use nn.ModuleList instead of nn.ModuleDict.