Saving custom layer parameters

[Yoshi]

@ananya

the track_grad_norm

Thank you! Let me try the flag.
Does the flag generate a json file or something?