@Evan_Vogelbaum How can I get all the flattened model parameters as numpy 1-D array or column vector through flatten_params() function? As @milad mentioned, after some computations on the weights, I want to unflatten the modified weights or replace the modified weights in the original model. Kindly help. Thank you.