### Check for existing issues - [x] Completed ### Description We should also save the weights and dimensions of the normalizing flow models.
Check for existing issues
Description
We should also save the weights and dimensions of the normalizing flow models.