Fix a mypy issue related to torch optimizer step (#3018)

* Fix a mypy issue

* Fix a flake issue in test_lr_finder

ignite/distributed/auto.py

@@ -1,5 +1,5 @@
 import warnings
-from typing import Any, Callable, Iterator, List, Optional, Union
+from typing import Any, Iterator, List, Optional, Union
 
 import torch
 import torch.nn as nn
@@ -351,5 +351,5 @@ def __init__(self, optimizer: Optimizer) -> None:
             super(self.__class__, self).__init__(optimizer.param_groups)  # type: ignore[call-arg]
             self.wrapped_optimizer = optimizer
 
-        def step(self, closure: Optional[Callable] = None) -> None:
+        def step(self, closure: Any = None) -> Any:
             xm.optimizer_step(self.wrapped_optimizer, barrier=True)

tests/ignite/handlers/test_lr_finder.py

@@ -405,13 +405,13 @@ def test_engine_output_type(lr_finder, dummy_engine, optimizer):
     lr_finder._history = {"lr": [], "loss": []}
     lr_finder._log_lr_and_loss(dummy_engine, output_transform=lambda x: x, smooth_f=0, diverge_th=1)
     loss = lr_finder._history["loss"][-1]
-    assert type(loss) == float
+    assert type(loss) is float
 
     dummy_engine.state.output = torch.tensor([10.0], dtype=torch.float32)
     lr_finder._history = {"lr": [], "loss": []}
     lr_finder._log_lr_and_loss(dummy_engine, output_transform=lambda x: x, smooth_f=0, diverge_th=1)
     loss = lr_finder._history["loss"][-1]
-    assert type(loss) == float
+    assert type(loss) is float
 
 
 def test_lr_suggestion_unexpected_curve(lr_finder, to_save, dummy_engine, dataloader):