Solved by nutszebra in post #2 
              
              
                If I were you, I set theta by register_buffer in the custom mlp like running_mean in batch norm (here is the link of that part). 
In this case, it’s not necessary to overwrite load_state_dict and state_dict.