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.6 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.

torch.save(modules, ‘modules.pt’)
torch.save(x1_transform, ‘x1_transform.pt’)
torch.save(modules.state_dict(), ‘modules_state_dict.pt’)

modules = torch.load(‘modules.pt’)
x1_transform = torch.load(‘x1_transform.pt’)
modules.load_state_dict(torch.load(“modules_state_dict.pt”))

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 ?