Fix mypy error

Author: sadra-barikbin

ignite/handlers/checkpoint.py

@@ -16,6 +16,10 @@
 if Version(torch.__version__) >= Version("1.9.0"):
     from torch.distributed.optim import ZeroRedundancyOptimizer
 
+    HAVE_ZERO = True
+else:
+    HAVE_ZERO = False
+
 import ignite.distributed as idist
 from ignite.base import Serializable
 from ignite.engine import Engine, Events
@@ -287,7 +291,7 @@ def __init__(
         filename_pattern: Optional[str] = None,
         include_self: bool = False,
         greater_or_equal: bool = False,
-        save_on_rank: Optional[int] = 0,
+        save_on_rank: int = 0,
     ):
 
         if not isinstance(to_save, collections.Mapping):
@@ -471,7 +475,7 @@ def _setup_checkpoint(self) -> Dict[str, Dict[Any, Any]]:
             for k, obj in self.to_save.items():
                 if isinstance(obj, (nn.DataParallel, nn.parallel.DistributedDataParallel)):
                     obj = obj.module
-                elif Version(torch.__version__) >= Version("1.9.0") and isinstance(obj, ZeroRedundancyOptimizer):
+                elif HAVE_ZERO and isinstance(obj, ZeroRedundancyOptimizer):
                     obj.consolidate_state_dict(to=self.save_on_rank)
                     if self.save_on_rank != idist.get_rank():
                         continue
@@ -791,7 +795,7 @@ def __init__(
         atomic: bool = True,
         create_dir: bool = True,
         require_empty: bool = True,
-        save_on_rank: Optional[int] = 0,
+        save_on_rank: int = 0,
         **kwargs: Any,
     ):
         self.dirname = Path(dirname).expanduser()
@@ -957,7 +961,7 @@ def __init__(
         filename_pattern: Optional[str] = None,
         include_self: bool = False,
         greater_or_equal: bool = False,
-        save_on_rank: Optional[int] = 0,
+        save_on_rank: int = 0,
         **kwargs: Any,
     ):
 

tests/ignite/handlers/test_checkpoint.py

@@ -11,9 +11,6 @@
 import torch.nn as nn
 from packaging.version import Version
 
-if Version(torch.__version__) >= Version("1.9.0"):
-    from torch.distributed.optim import ZeroRedundancyOptimizer
-
 import ignite.distributed as idist
 from ignite.engine import Engine, Events, State
 from ignite.handlers import Checkpoint, DiskSaver, EarlyStopping, global_step_from_engine, ModelCheckpoint
@@ -1247,6 +1244,9 @@ def _test_checkpoint_load_objects_ddp(device):
 
 
 def _test_checkpoint_with_ZeRO(device, dirname, local_rank):
+
+    from torch.distributed.optim import ZeroRedundancyOptimizer
+
     model = DummyModel().to(device)
     opt = ZeroRedundancyOptimizer(model.parameters(), torch.optim.SGD, lr=0.01)
     mocked_opt = MagicMock(ZeroRedundancyOptimizer, wraps=opt)
@@ -1282,7 +1282,9 @@ def test_distrib_gloo_cpu_or_gpu(distributed_context_single_node_gloo, dirname,
     _test_checkpoint_with_ddp(device)
     _test_checkpoint_load_objects_ddp(device)
 
-    if Version(torch.__version__) >= Version("1.9.0"):
+    from ignite.handlers.checkpoint import HAVE_ZERO
+
+    if HAVE_ZERO:
         _test_checkpoint_with_ZeRO(device, dirname, local_rank)