Hey @HaydenFaulkner, I would suggest trying to rewrite your model in Gluon, that would make it trivial to update the parameters and have better control on the flow of your data. Happy to assist you converting your Module code to Gluon, and hopefully convert you from Module to being a Gluon fan 
If that’s not an option, feel free to post your code here