Avoid unpickling the extra state when not needed#3123

ksivaman merged 11 commits into

NVIDIA:mainNVIDIA/TransformerEngine:mainfrom

ptrendx:pr_avoid_unpickleptrendx/TransformerEngine:pr_avoid_unpickleCopy head branch name to clipboard