How to flatten and then unflatten _all_ model parameters

@milad Did you get the solution? If yes, please share. I am also having the same question. Thank you.