Add missing reference of state dict hook (#4764)

faa6bebe · Yuge Zhang · GitHub · e0ce406f · faa6bebe
Unverified Commit faa6bebe authored Apr 14, 2022 by Yuge Zhang Committed by GitHub Apr 14, 2022
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

nni/retiarii/utils.py nni/retiarii/utils.py +1 -1

No files found.
--- a/nni/retiarii/utils.py
+++ b/nni/retiarii/utils.py
@@ -11,7 +11,7 @@ from pathlib import Path

 from nni.common.hpo_utils import ParameterSpec

-__all__ = ['NoContextError', 'ContextStack', 'ModelNamespace']
+__all__ = ['NoContextError', 'ContextStack', 'ModelNamespace', 'original_state_dict_hooks']


 def import_(target: str, allow_none: bool = False) -> Any: