[ONNX] Reland: Update training state logic to support ScriptedModule (#86745)
In https://github.com/pytorch/pytorch/issues/86325, it was reported that ScriptedModule do not have a training attribute and will fail export because we don't expect them as input.
Also
- Parameterized the test_util_funs test
Thanks @borisfom for the suggestion!
Fixes #86325
Pull Request resolved: https://github.com/pytorch/pytorch/pull/86745
Approved by: https://github.com/AllenTiTaiWang, https://github.com/BowenBao
