pytorch-ignite 0.6.0.dev20250310__py3-none-any.whl → 0.6.0.dev20260101__py3-none-any.whl

ignite/distributed/utils.py

@@ -2,10 +2,9 @@ import itertools
 import socket
 from contextlib import contextmanager
 from functools import wraps
-from typing import Any, Callable, List, Mapping, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, cast, List, Mapping, Optional, Sequence, Tuple, Union
 
 import torch
-from torch import distributed as dist
 
 from ignite.distributed.comp_models import (
     _SerialModel,
@@ -384,7 +383,7 @@ def all_gather_tensors_with_shapes(
     if isinstance(group, list) and all(isinstance(item, int) for item in group):
         group = _model.new_group(group)
 
-    if isinstance(_model, _SerialModel) or group == dist.GroupMember.NON_GROUP_MEMBER:
+    if _rank_not_in_group(group):
         return [tensor]
 
     max_shape = torch.tensor(shapes).amax(dim=0)
@@ -392,7 +391,7 @@ def all_gather_tensors_with_shapes(
     padded_tensor = torch.nn.functional.pad(
         tensor, tuple(itertools.chain.from_iterable(map(lambda dim_size: (0, dim_size), reversed(padding_sizes))))
     )
-    all_padded_tensors: torch.Tensor = _model.all_gather(padded_tensor, group=group)
+    all_padded_tensors: torch.Tensor = cast(torch.Tensor, _model.all_gather(padded_tensor, group=group))
     return [
         all_padded_tensors[
             [
@@ -731,3 +730,12 @@ def one_rank_first(rank: int = 0, local: bool = False) -> Any:
 
     if current_rank == rank:
         barrier()
+
+
+def _rank_not_in_group(group: Optional[Union[Any, List[int]]]) -> bool:
+    """Check if the current process's rank is not in a given group."""
+    if group is None:
+        return False
+    if isinstance(group, list) and all(isinstance(item, int) for item in group):
+        group = new_group(group)
+    return _model._rank_not_in_group(group)

ignite/engine/__init__.py

@@ -133,11 +133,11 @@ def supervised_training_step_amp(
     prepare_batch: Callable = _prepare_batch,
     model_transform: Callable[[Any], Any] = lambda output: output,
     output_transform: Callable[[Any, Any, Any, torch.Tensor], Any] = lambda x, y, y_pred, loss: loss.item(),
-    scaler: Optional["torch.cuda.amp.GradScaler"] = None,
+    scaler: Optional["torch.amp.GradScaler"] = None,
     gradient_accumulation_steps: int = 1,
     model_fn: Callable[[torch.nn.Module, Any], Any] = lambda model, x: model(x),
 ) -> Callable:
-    """Factory function for supervised training using ``torch.cuda.amp``.
+    """Factory function for supervised training using ``torch.amp``.
 
     Args:
         model: the model to train.
@@ -170,7 +170,7 @@ def supervised_training_step_amp(
             model = ...
             optimizer = ...
             loss_fn = ...
-            scaler = torch.cuda.amp.GradScaler(2**10)
+            scaler = torch.amp.GradScaler('cuda', 2**10)
 
             update_fn = supervised_training_step_amp(model, optimizer, loss_fn, 'cuda', scaler=scaler)
             trainer = Engine(update_fn)
@@ -393,8 +393,8 @@ def supervised_training_step_tpu(
 
 
 def _check_arg(
-    on_tpu: bool, on_mps: bool, amp_mode: Optional[str], scaler: Optional[Union[bool, "torch.cuda.amp.GradScaler"]]
-) -> Tuple[Optional[str], Optional["torch.cuda.amp.GradScaler"]]:
+    on_tpu: bool, on_mps: bool, amp_mode: Optional[str], scaler: Optional[Union[bool, "torch.amp.GradScaler"]]
+) -> Tuple[Optional[str], Optional["torch.amp.GradScaler"]]:
     """Checking tpu, mps, amp and GradScaler instance combinations."""
     if on_mps and amp_mode:
         raise ValueError("amp_mode cannot be used with mps device. Consider using amp_mode=None or device='cuda'.")
@@ -410,9 +410,9 @@ def _check_arg(
             raise ValueError(f"scaler argument is {scaler}, but amp_mode is {amp_mode}. Consider using amp_mode='amp'.")
         elif amp_mode == "amp" and isinstance(scaler, bool):
             try:
-                from torch.cuda.amp import GradScaler
+                from torch.amp import GradScaler
             except ImportError:
-                raise ImportError("Please install torch>=1.6.0 to use scaler argument.")
+                raise ImportError("Please install torch>=2.3.1 to use scaler argument.")
             scaler = GradScaler(enabled=True)
 
     if on_tpu:
@@ -434,7 +434,7 @@ def create_supervised_trainer(
     output_transform: Callable[[Any, Any, Any, torch.Tensor], Any] = lambda x, y, y_pred, loss: loss.item(),
     deterministic: bool = False,
     amp_mode: Optional[str] = None,
-    scaler: Union[bool, "torch.cuda.amp.GradScaler"] = False,
+    scaler: Union[bool, "torch.amp.GradScaler"] = False,
     gradient_accumulation_steps: int = 1,
     model_fn: Callable[[torch.nn.Module, Any], Any] = lambda model, x: model(x),
 ) -> Engine:
@@ -459,7 +459,7 @@ def create_supervised_trainer(
             :class:`~ignite.engine.deterministic.DeterministicEngine`, otherwise :class:`~ignite.engine.engine.Engine`
             (default: False).
         amp_mode: can be ``amp`` or ``apex``, model and optimizer will be casted to float16 using
-            `torch.cuda.amp <https://pytorch.org/docs/stable/amp.html>`_ for ``amp`` and
+            `torch.amp <https://pytorch.org/docs/stable/amp.html>`_ for ``amp`` and
             using `apex <https://nvidia.github.io/apex>`_ for ``apex``. (default: None)
         scaler: GradScaler instance for gradient scaling if `torch>=1.6.0`
             and ``amp_mode`` is ``amp``. If ``amp_mode`` is ``apex``, this argument will be ignored.

ignite/engine/deterministic.py

@@ -85,7 +85,7 @@ class ReproducibleBatchSampler(BatchSampler):
 
 
 def _get_rng_states() -> List[Any]:
-    output = [random.getstate(), torch.get_rng_state()]
+    output: List[Any] = [random.getstate(), torch.get_rng_state()]
     try:
         import numpy as np
 

ignite/engine/engine.py

@@ -148,7 +148,7 @@ class Engine(Serializable):
         self.should_interrupt = False
         self.state = State()
         self._state_dict_user_keys: List[str] = []
-        self._allowed_events: List[EventEnum] = []
+        self._allowed_events: List[Union[str, EventEnum]] = []
 
         self._dataloader_iter: Optional[Iterator[Any]] = None
         self._init_iter: Optional[int] = None
@@ -163,9 +163,7 @@ class Engine(Serializable):
         # generator provided by self._internal_run_as_gen
         self._internal_run_generator: Optional[Generator[Any, None, State]] = None
 
-    def register_events(
-        self, *event_names: Union[List[str], List[EventEnum]], event_to_attr: Optional[dict] = None
-    ) -> None:
+    def register_events(self, *event_names: Union[str, EventEnum], event_to_attr: Optional[dict] = None) -> None:
         """Add events that can be fired.
 
         Registering an event will let the user trigger these events at any point.
@@ -249,6 +247,17 @@ class Engine(Serializable):
         # we need to update state attributes associated with new custom events
         self.state._update_attrs()
 
+    def has_registered_events(self, event: Any) -> bool:
+        """Check whether engine has a registered event.
+
+        Args:
+            event: Event to check for registration.
+
+        Returns:
+            bool: True if the event is registered, False otherwise.
+        """
+        return event in self._allowed_events
+
     def _handler_wrapper(self, handler: Callable, event_name: Any, event_filter: Callable) -> Callable:
         # signature of the following wrapper will be inspected during registering to check if engine is necessary
         # we have to build a wrapper with relevant signature : solution is functools.wraps
@@ -328,7 +337,7 @@ class Engine(Serializable):
 
         try:
             _check_signature(handler, "handler", self, *(event_args + args), **kwargs)
-            self._event_handlers[event_name].append((handler, (self,) + args, kwargs))
+            self._event_handlers[event_name].append((handler, (weakref.ref(self),) + args, kwargs))
         except ValueError:
             _check_signature(handler, "handler", *(event_args + args), **kwargs)
             self._event_handlers[event_name].append((handler, args, kwargs))
@@ -432,7 +441,15 @@ class Engine(Serializable):
         self.last_event_name = event_name
         for func, args, kwargs in self._event_handlers[event_name]:
             kwargs.update(event_kwargs)
-            first, others = ((args[0],), args[1:]) if (args and args[0] == self) else ((), args)
+            if args and isinstance(args[0], weakref.ref):
+                resolved_engine = args[0]()
+                if resolved_engine is None:
+                    raise RuntimeError("Engine reference not resolved. Cannot execute event handler.")
+                first, others = ((resolved_engine,), args[1:])
+            else:
+                # metrics do not provide engine when registered
+                first, others = (tuple(), args)
+
             func(*first, *(event_args + others), **kwargs)
 
     def fire_event(self, event_name: Any) -> None:
@@ -970,9 +987,9 @@ class Engine(Serializable):
     def _internal_run_as_gen(self) -> Generator[Any, None, State]:
         self.should_terminate = self.should_terminate_single_epoch = self.should_interrupt = False
         self._init_timers(self.state)
+        start_time = time.time()
         try:
             try:
-                start_time = time.time()
                 self._fire_event(Events.STARTED)
                 yield from self._maybe_terminate_or_interrupt()
 
@@ -991,7 +1008,7 @@ class Engine(Serializable):
                     # time is available for handlers but must be updated after fire
                     self.state.times[Events.EPOCH_COMPLETED.name] = epoch_time_taken
 
-                    if self.should_terminate_single_epoch != "skip_epoch_completed":  # type: ignore[comparison-overlap]
+                    if self.should_terminate_single_epoch != "skip_epoch_completed":
                         handlers_start_time = time.time()
                         self._fire_event(Events.EPOCH_COMPLETED)
                         epoch_time_taken += time.time() - handlers_start_time
@@ -1024,7 +1041,7 @@ class Engine(Serializable):
             self.state.times[Events.COMPLETED.name] = time_taken
 
             # do not fire Events.COMPLETED if we terminated the run with flag `skip_completed=True`
-            if self.should_terminate != "skip_completed":  # type: ignore[comparison-overlap]
+            if self.should_terminate != "skip_completed":
                 handlers_start_time = time.time()
                 self._fire_event(Events.COMPLETED)
                 time_taken += time.time() - handlers_start_time
@@ -1069,7 +1086,7 @@ class Engine(Serializable):
                 )
 
             while True:
-                self.state.batch = self.state.output = None
+                self.state.batch = None
 
                 try:
                     # Avoid Events.GET_BATCH_STARTED triggered twice when data iter is restarted
@@ -1081,6 +1098,9 @@ class Engine(Serializable):
                             yield from self._maybe_terminate_or_interrupt()
 
                     self.state.batch = next(self._dataloader_iter)
+                    # We on purpose reset state.output here as for iterable dataloaders
+                    # we accidentally can remove it when one epoch is completed.
+                    self.state.output = None
 
                     # We should not trigger GET_BATCH_STARTED, GET_BATCH_COMPLETED, DATALOADER_STOP_ITERATION events
                     # if no data was provided to engine.run(data=None, ...)
@@ -1167,9 +1187,9 @@ class Engine(Serializable):
         # internal_run without generator for BC
         self.should_terminate = self.should_terminate_single_epoch = self.should_interrupt = False
         self._init_timers(self.state)
+        start_time = time.time()
         try:
             try:
-                start_time = time.time()
                 self._fire_event(Events.STARTED)
                 self._maybe_terminate_legacy()
 
@@ -1188,7 +1208,7 @@ class Engine(Serializable):
                     # time is available for handlers but must be updated after fire
                     self.state.times[Events.EPOCH_COMPLETED.name] = epoch_time_taken
 
-                    if self.should_terminate_single_epoch != "skip_epoch_completed":  # type: ignore[comparison-overlap]
+                    if self.should_terminate_single_epoch != "skip_epoch_completed":
                         handlers_start_time = time.time()
                         self._fire_event(Events.EPOCH_COMPLETED)
                         epoch_time_taken += time.time() - handlers_start_time
@@ -1221,7 +1241,7 @@ class Engine(Serializable):
             self.state.times[Events.COMPLETED.name] = time_taken
 
             # do not fire Events.COMPLETED if we terminated the run with flag `skip_completed=True`
-            if self.should_terminate != "skip_completed":  # type: ignore[comparison-overlap]
+            if self.should_terminate != "skip_completed":
                 handlers_start_time = time.time()
                 self._fire_event(Events.COMPLETED)
                 time_taken += time.time() - handlers_start_time
@@ -1254,7 +1274,7 @@ class Engine(Serializable):
                 )
 
             while True:
-                self.state.batch = self.state.output = None
+                self.state.batch = None
                 try:
                     # Avoid Events.GET_BATCH_STARTED triggered twice when data iter is restarted
                     if self.last_event_name != Events.DATALOADER_STOP_ITERATION:
@@ -1265,6 +1285,10 @@ class Engine(Serializable):
                             self._maybe_terminate_legacy()
 
                     self.state.batch = next(self._dataloader_iter)
+                    # We on purpose reset state.output here as for iterable dataloaders
+                    # we accidentally can remove it when one epoch is completed.
+                    self.state.output = None
+
                     # We should not trigger GET_BATCH_STARTED, GET_BATCH_COMPLETED, DATALOADER_STOP_ITERATION events
                     # if no data was provided to engine.run(data=None, ...)
                     if self.state.dataloader is not None:

ignite/engine/events.py

@@ -91,7 +91,7 @@ class CallableEventWithFilter:
             raise ValueError("Argument every should be integer and greater than zero")
 
         if once is not None:
-            c1 = isinstance(once, numbers.Integral) and once > 0
+            c1 = isinstance(once, int) and once > 0
             c2 = isinstance(once, Sequence) and len(once) > 0 and all(isinstance(e, int) and e > 0 for e in once)
             if not (c1 or c2):
                 raise ValueError(
@@ -240,6 +240,7 @@ class EventEnum(CallableEventWithFilter, Enum):
     def __new__(cls, value: str) -> "EventEnum":
         obj = CallableEventWithFilter.__new__(cls)
         obj._value_ = value
+        # pyrefly: ignore [bad-return]
         return obj
 
 

ignite/handlers/__init__.py

@@ -6,6 +6,7 @@ from ignite.handlers.checkpoint import Checkpoint, DiskSaver, ModelCheckpoint
 from ignite.handlers.clearml_logger import ClearMLLogger
 from ignite.handlers.early_stopping import EarlyStopping
 from ignite.handlers.ema_handler import EMAHandler
+from ignite.handlers.fbresearch_logger import FBResearchLogger
 from ignite.handlers.lr_finder import FastaiLRFinder
 from ignite.handlers.mlflow_logger import MLflowLogger
 from ignite.handlers.neptune_logger import NeptuneLogger
@@ -64,6 +65,7 @@ __all__ = [
     "CyclicalScheduler",
     "create_lr_scheduler_with_warmup",
     "FastaiLRFinder",
+    "FBResearchLogger",
     "EMAHandler",
     "BasicTimeProfiler",
     "HandlersTimeProfiler",

ignite/handlers/base_logger.py

@@ -1,5 +1,6 @@
 """Base logger and its helper handlers."""
 
+import collections.abc as collections
 import numbers
 import warnings
 from abc import ABCMeta, abstractmethod
@@ -145,30 +146,64 @@ class BaseOutputHandler(BaseHandler):
 
         metrics_state_attrs_dict: Dict[Any, Union[str, float, numbers.Number]] = OrderedDict()
 
-        def key_tuple_tf(tag: str, name: str, *args: str) -> Tuple[str, ...]:
-            return (tag, name) + args
+        def key_tuple_fn(parent_key: Union[str, Tuple[str, ...]], *args: str) -> Tuple[str, ...]:
+            if parent_key is None or isinstance(parent_key, str):
+                return (parent_key,) + args
+            return parent_key + args
 
-        def key_str_tf(tag: str, name: str, *args: str) -> str:
-            return "/".join((tag, name) + args)
+        def key_str_fn(parent_key: str, *args: str) -> str:
+            args_str = "/".join(args)
+            return f"{parent_key}/{args_str}"
 
-        key_tf = key_tuple_tf if key_tuple else key_str_tf
+        key_fn = key_tuple_fn if key_tuple else key_str_fn
 
-        for name, value in metrics_state_attrs.items():
+        def handle_value_fn(
+            value: Union[str, int, float, numbers.Number, torch.Tensor]
+        ) -> Union[None, str, float, numbers.Number]:
             if isinstance(value, numbers.Number):
-                metrics_state_attrs_dict[key_tf(self.tag, name)] = value
+                return value
             elif isinstance(value, torch.Tensor) and value.ndimension() == 0:
-                metrics_state_attrs_dict[key_tf(self.tag, name)] = value.item()
-            elif isinstance(value, torch.Tensor) and value.ndimension() == 1:
-                for i, v in enumerate(value):
-                    metrics_state_attrs_dict[key_tf(self.tag, name, str(i))] = v.item()
+                return value.item()
             else:
                 if isinstance(value, str) and log_text:
-                    metrics_state_attrs_dict[key_tf(self.tag, name)] = value
+                    return value
                 else:
                     warnings.warn(f"Logger output_handler can not log metrics value type {type(value)}")
+            return None
+
+        metrics_state_attrs_dict = _flatten_dict(metrics_state_attrs, key_fn, handle_value_fn, parent_key=self.tag)
         return metrics_state_attrs_dict
 
 
+def _flatten_dict(
+    in_dict: collections.Mapping,
+    key_fn: Callable,
+    value_fn: Callable,
+    parent_key: Optional[Union[str, Tuple[str, ...]]] = None,
+) -> Dict:
+    items = {}
+    for key, value in in_dict.items():
+        new_key = key_fn(parent_key, key)
+        if isinstance(value, collections.Mapping):
+            items.update(_flatten_dict(value, key_fn, value_fn, new_key))
+        elif any(
+            [
+                isinstance(value, tuple) and hasattr(value, "_fields"),  # namedtuple
+                not isinstance(value, str) and isinstance(value, collections.Sequence),
+            ]
+        ):
+            for i, item in enumerate(value):
+                items.update(_flatten_dict({str(i): item}, key_fn, value_fn, new_key))
+        elif isinstance(value, torch.Tensor) and value.ndimension() == 1:
+            for i, item in enumerate(value):
+                items.update(_flatten_dict({str(i): item.item()}, key_fn, value_fn, new_key))
+        else:
+            new_value = value_fn(value)
+            if new_value is not None:
+                items[new_key] = new_value
+    return items
+
+
 class BaseWeightsScalarHandler(BaseWeightsHandler):
     """
     Helper handler to log model's weights or gradients as scalars.

ignite/handlers/checkpoint.py

@@ -21,10 +21,21 @@ else:
 
 import ignite.distributed as idist
 from ignite.base import Serializable
-from ignite.engine import Engine, Events
+from ignite.engine import Engine, Events, EventEnum
 from ignite.utils import _tree_apply2, _tree_map
 
-__all__ = ["Checkpoint", "DiskSaver", "ModelCheckpoint", "BaseSaveHandler"]
+__all__ = ["Checkpoint", "DiskSaver", "ModelCheckpoint", "BaseSaveHandler", "CheckpointEvents"]
+
+
+class CheckpointEvents(EventEnum):
+    """Events fired by :class:`~ignite.handlers.checkpoint.Checkpoint`
+
+    - SAVED_CHECKPOINT : triggered when checkpoint handler has saved objects
+
+    .. versionadded:: 0.5.3
+    """
+
+    SAVED_CHECKPOINT = "saved_checkpoint"
 
 
 class BaseSaveHandler(metaclass=ABCMeta):
@@ -264,6 +275,29 @@ class Checkpoint(Serializable):
                 to_save, save_handler=DiskSaver('/tmp/models', create_dir=True, **kwargs), n_saved=2
             )
 
+        Respond to checkpoint events:
+
+        .. code-block:: python
+
+            from ignite.handlers import Checkpoint
+            from ignite.engine import Engine, Events
+
+            checkpoint_handler = Checkpoint(
+                {'model': model, 'optimizer': optimizer},
+                save_dir,
+                n_saved=2
+            )
+
+            @trainer.on(Checkpoint.SAVED_CHECKPOINT)
+            def on_checkpoint_saved(engine):
+                print(f"Checkpoint saved at epoch {engine.state.epoch}")
+
+            trainer.add_event_handler(Events.EPOCH_COMPLETED, checkpoint_handler)
+
+    Attributes:
+        SAVED_CHECKPOINT: Alias of ``SAVED_CHECKPOINT`` from
+            :class:`~ignite.handlers.checkpoint.CheckpointEvents`.
+
     .. versionchanged:: 0.4.3
 
         - Checkpoint can save model with same filename.
@@ -274,9 +308,14 @@ class Checkpoint(Serializable):
         - `score_name` can be used to define `score_function` automatically without providing `score_function`.
         - `save_handler` automatically saves to disk if path to directory is provided.
         - `save_on_rank` saves objects on this rank in a distributed configuration.
+
+    .. versionchanged:: 0.5.3
+
+        - Added ``SAVED_CHECKPOINT`` class attribute.
     """
 
-    Item = NamedTuple("Item", [("priority", int), ("filename", str)])
+    SAVED_CHECKPOINT = CheckpointEvents.SAVED_CHECKPOINT
+    Item = NamedTuple("Item", [("priority", Union[int, float]), ("filename", str)])
     _state_dict_all_req_keys = ("_saved",)
 
     def __init__(
@@ -284,7 +323,7 @@ class Checkpoint(Serializable):
         to_save: Mapping,
         save_handler: Union[str, Path, Callable, BaseSaveHandler],
         filename_prefix: str = "",
-        score_function: Optional[Callable] = None,
+        score_function: Optional[Callable[[Engine], Union[int, float]]] = None,
         score_name: Optional[str] = None,
         n_saved: Union[int, None] = 1,
         global_step_transform: Optional[Callable] = None,
@@ -400,6 +439,8 @@ class Checkpoint(Serializable):
             return new > self._saved[0].priority
 
     def __call__(self, engine: Engine) -> None:
+        if not engine.has_registered_events(CheckpointEvents.SAVED_CHECKPOINT):
+            engine.register_events(*CheckpointEvents)
         global_step = None
         if self.global_step_transform is not None:
             global_step = self.global_step_transform(engine, engine.last_event_name)
@@ -460,11 +501,11 @@ class Checkpoint(Serializable):
             if self.include_self:
                 # Now that we've updated _saved, we can add our own state_dict.
                 checkpoint["checkpointer"] = self.state_dict()
-
             try:
                 self.save_handler(checkpoint, filename, metadata)
             except TypeError:
                 self.save_handler(checkpoint, filename)
+            engine.fire_event(CheckpointEvents.SAVED_CHECKPOINT)
 
     def _setup_checkpoint(self) -> Dict[str, Any]:
         if self.to_save is not None:

ignite/handlers/clearml_logger.py

@@ -109,8 +109,17 @@ class ClearMLLogger(BaseLogger):
                 log_handler=WeightsScalarHandler(model)
             )
 
+    Note:
+        :class:`~ignite.handlers.clearml_logger.OutputHandler` can handle
+        metrics, state attributes and engine output values of the following format:
+        - scalar values (i.e. int, float)
+        - 0d and 1d pytorch tensors
+        - dicts and list/tuples of previous types
+
     """
 
+    _task: Any
+
     def __init__(self, **kwargs: Any):
         try:
             from clearml import Task
@@ -342,9 +351,10 @@ class OutputHandler(BaseOutputHandler):
         for key, value in metrics.items():
             if len(key) == 2:
                 logger.clearml_logger.report_scalar(title=key[0], series=key[1], iteration=global_step, value=value)
-            elif len(key) == 3:
+            elif len(key) >= 3:
+                series = "/".join(key[2:])
                 logger.clearml_logger.report_scalar(
-                    title=f"{key[0]}/{key[1]}", series=key[2], iteration=global_step, value=value
+                    title=f"{key[0]}/{key[1]}", series=series, iteration=global_step, value=value
                 )
 
 
@@ -815,6 +825,8 @@ class ClearMLSaver(DiskSaver):
 
     """
 
+    _task: Any
+
     def __init__(
         self,
         logger: Optional[ClearMLLogger] = None,
@@ -949,8 +961,8 @@ class ClearMLSaver(DiskSaver):
             metadata=metadata,
         )
 
-        pre_cb_id = WeightsFileHandler.add_pre_callback(cb_context.pre_callback)
-        post_cb_id = WeightsFileHandler.add_post_callback(cb_context.post_callback)
+        pre_cb_id = WeightsFileHandler.add_pre_callback(cb_context.pre_callback)  # type: ignore[arg-type]
+        post_cb_id = WeightsFileHandler.add_post_callback(cb_context.post_callback)  # type: ignore[arg-type]
 
         try:
             super(ClearMLSaver, self).__call__(checkpoint, filename, metadata)

ignite/handlers/fbresearch_logger.py

@@ -7,7 +7,7 @@ import torch
 
 from ignite import utils
 from ignite.engine import Engine, Events
-from ignite.handlers import Timer
+from ignite.handlers.timing import Timer
 
 MB = 1024.0 * 1024.0
 
@@ -154,7 +154,7 @@ class FBResearchLogger:
         if torch.cuda.is_available():
             cuda_max_mem = f"GPU Max Mem: {torch.cuda.max_memory_allocated() / MB:.0f} MB"
 
-        current_iter = engine.state.iteration % (engine.state.epoch_length + 1)
+        current_iter = ((engine.state.iteration - 1) % engine.state.epoch_length) + 1
         iter_avg_time = self.iter_timer.value()
 
         eta_seconds = iter_avg_time * (engine.state.epoch_length - current_iter)

ignite/handlers/lr_finder.py

@@ -98,15 +98,18 @@ class FastaiLRFinder:
         self._best_loss = None
         self._diverge_flag = False
 
+        assert trainer.state.epoch_length is not None
+        assert trainer.state.max_epochs is not None
+
         # attach LRScheduler to trainer.
         if num_iter is None:
             num_iter = trainer.state.epoch_length * trainer.state.max_epochs
         else:
-            max_iter = trainer.state.epoch_length * trainer.state.max_epochs  # type: ignore[operator]
+            max_iter = trainer.state.epoch_length * trainer.state.max_epochs
             if max_iter < num_iter:
                 max_iter = num_iter
                 trainer.state.max_iters = num_iter
-                trainer.state.max_epochs = ceil(num_iter / trainer.state.epoch_length)  # type: ignore[operator]
+                trainer.state.max_epochs = ceil(num_iter / trainer.state.epoch_length)
 
         if not trainer.has_event_handler(self._reached_num_iterations):
             trainer.add_event_handler(Events.ITERATION_COMPLETED, self._reached_num_iterations, num_iter)
@@ -178,17 +181,14 @@ class FastaiLRFinder:
         loss = idist.all_reduce(loss)
         lr = self._lr_schedule.get_param()
         self._history["lr"].append(lr)
-        if trainer.state.iteration == 1:
+        if trainer.state.iteration != 1 and smooth_f > 0:
+            loss = smooth_f * loss + (1 - smooth_f) * self._history["loss"][-1]
+        if self._best_loss is None or loss < self._best_loss:
             self._best_loss = loss
-        else:
-            if smooth_f > 0:
-                loss = smooth_f * loss + (1 - smooth_f) * self._history["loss"][-1]
-            if loss < self._best_loss:
-                self._best_loss = loss
         self._history["loss"].append(loss)
 
         # Check if the loss has diverged; if it has, stop the trainer
-        if self._history["loss"][-1] > diverge_th * self._best_loss:  # type: ignore[operator]
+        if self._history["loss"][-1] > diverge_th * self._best_loss:
             self._diverge_flag = True
             self.logger.info("Stopping early, the loss has diverged")
             trainer.terminate()

ignite/handlers/mlflow_logger.py

@@ -84,6 +84,49 @@ class MLflowLogger(BaseLogger):
                 optimizer=optimizer,
                 param_name='lr'  # optional
             )
+
+    Note:
+        :class:`~ignite.handlers.mlflow_logger.OutputHandler` can handle
+        metrics, state attributes and engine output values of the following format:
+        - scalar values (i.e. int, float)
+        - 0d and 1d pytorch tensors
+        - dicts and list/tuples of previous types
+
+        .. code-block:: python
+
+            # !!! This is not a runnable code !!!
+            evalutator.state.metrics = {
+                "a": 0,
+                "dict_value": {
+                    "a": 111,
+                    "c": {"d": 23, "e": [123, 234]},
+                },
+                "list_value": [12, 13, {"aa": 33, "bb": 44}],
+                "tuple_value": (112, 113, {"aaa": 33, "bbb": 44}),
+            }
+
+            handler = OutputHandler(
+                tag="tag",
+                metric_names="all",
+            )
+
+            handler(evaluator, mlflow_logger, event_name=Events.EPOCH_COMPLETED)
+            # Behind it would call `mlflow_logger.log_metrics` on
+            # {
+            #     "tag/a": 0,
+            #     "tag/dict_value/a": 111,
+            #     "tag/dict_value/c/d": 23,
+            #     "tag/dict_value/c/e/0": 123,
+            #     "tag/dict_value/c/e/1": 234,
+            #     "tag/list_value/0": 12,
+            #     "tag/list_value/1": 13,
+            #     "tag/list_value/2/aa": 33,
+            #     "tag/list_value/2/bb": 44,
+            #     "tag/tuple_value/0": 112,
+            #     "tag/tuple_value/1": 113,
+            #     "tag/tuple_value/2/aaa": 33,
+            #     "tag/tuple_value/2/bbb": 44,
+            # }
     """
 
     def __init__(self, tracking_uri: Optional[str] = None):