How to save and load Normalizing Flows Model and distributions

Following the steps in the Normalizing Flows tutorial, I was able to train a conditional model of my own. Now I am trying to save the module and distributions following the documentation and tutorials (Deep Markov Model — Pyro Tutorials 1.8.0 documentation). I am able to save and load all the components separately using the following commands but the loaded model doesn’t seem to share the properties of the original model., ‘’), ‘’), ‘’)

modules = torch.load(‘’)
x1_transform = torch.load(‘’)

The saved model always results in a similar trace

The loaded model always results in this type of output:

However, plotting the distributions are similar before and after loading:

Saved model:
correct 2

Loaded Model:
faulty 2

I am very new to the pyro ecosystem and I would be grateful for your help.

Any ideas @stefanwebb ?