DMM example- Visualize gradients?

Hi!

First, thanks for your time and attention. Second, I am following the DMM model https://pyro.ai/examples/dmm.html to build my own model and I am struggling to understand how can I visualize the gradients.

I know in the Gaussian Mixture Model https://pyro.ai/examples/gmm.html?highlight=gaussian%20mixture%20model is easy to do:

 # Register hooks to monitor gradient norms.
 gradient_norms = defaultdict(list)
 for name, value in pyro.get_param_store().named_parameters():
     value.register_hook(lambda g, name=name: gradient_norms[name].append(g.norm().item()))

However in the DMM is not that direct. I need to initialize the DMM somehow, because right now I am obtaining an empty graph :frowning:

Best regards and thanks for any help :slight_smile: