Avoid module level StrictHandling

Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com>
NVIDIA · Jul 4, 2024 · ddb6be3 · ddb6be3
1 parent f6f4e9e
commit ddb6be3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nemo/utils/callbacks/dist_ckpt_io.py b/nemo/utils/callbacks/dist_ckpt_io.py
@@ -184,7 +184,7 @@ def __init__(
         self,
         save_ckpt_format: str,
         load_directly_on_device: bool = True,
-        load_strictness: StrictHandling = StrictHandling.ASSUME_OK_UNEXPECTED,
+        load_strictness: 'StrictHandling' = 'assume_ok_unexpected',
         async_save: bool = False,
     ):
         super().__init__()
@@ -245,7 +245,7 @@ def load_checkpoint(
         path: _PATH,
         map_location: Optional[Any] = None,
         sharded_state_dict: Dict[str, Any] = None,
-        strict: Union[None, bool, StrictHandling] = None,
+        strict: Union[None, bool, 'StrictHandling'] = None,
         validate_access_integrity: Optional[bool] = True,
     ) -> Dict[str, Any]:
         """Loads a distributed checkpoint.