code: def _test_save_model_optimizer_lr_scheduler_with_state_dict(
      device, on_zero_rank=False
  ):
      if idist.get_rank() == 0:
          clearml.Task.current_task = MagicMock(spec=clearml.Task)
          clearml.binding.frameworks.WeightsFileHandler.create_output_model = (
              MagicMock()
          )
  
      torch.manual_seed(23)
  
      model = DummyModel().to(device)
  
      optim = torch.optim.SGD(model.parameters(), lr=0.1)
      lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(
          optim, gamma=0.5
      )
  
      def update_fn(engine, batch):
          x = torch.rand((4, 2)).to(device)
          optim.zero_grad()
          y = model(x)
          # Below code raises: RuntimeError: torch_xla/csrc/tensor_impl.cpp:144 : XLA tensors do not have storage
          # Probably related to https://github.com/pytorch/xla/issues/2576
          # loss = y.pow(2.0).sum()
          loss = y.sum()
          loss.backward()
          if idist.has_xla_support:
              import torch_xla.core.xla_model as xm
  
              xm.optimizer_step(optim, barrier=True)
          else:
              optim.step()
          lr_scheduler.step()
  
      engine = Engine(update_fn)
  
      to_save = {
          "model": model,
          "optimizer": optim,
          "lr_scheduler": lr_scheduler,
      }
  
      with pytest.warns(
          UserWarning,
          match=r"ClearMLSaver created a temporary checkpoints directory",
      ):
          clearml_saver = ClearMLSaver()
  
      if (not on_zero_rank) or (on_zero_rank and idist.get_rank() == 0):
          checkpoint = Checkpoint(
              to_save=to_save, save_handler=clearml_saver, n_saved=1
          )
          engine.add_event_handler(Events.EPOCH_COMPLETED, checkpoint)
  
      engine.run([0], max_epochs=4)
  
      idist.barrier()
  
      saved_objects = sorted(os.listdir(clearml_saver.dirname))
      # saved object is ['PREFIX_checkpoint_3.pt', ]
      saved_checkpoint = clearml_saver.dirname / saved_objects[0]
  
      if idist.has_xla_support:
          device = "cpu"
  
      loaded_obj = torch.load(saved_checkpoint, map_location=device)
      for f in ["model", "optimizer", "lr_scheduler"]:
          assert f in loaded_obj
      loaded_model_state_dict = loaded_obj["model"]
      loaded_optimizer_state_dict = loaded_obj["optimizer"]
      loaded_lr_scheduler_state_dict = loaded_obj["lr_scheduler"]
  
      assert isinstance(loaded_model_state_dict, dict)
      assert isinstance(loaded_optimizer_state_dict, dict)
      assert isinstance(loaded_lr_scheduler_state_dict, dict)
  
      # Specifically move device to CPU first
      model_state_dict = model.cpu().state_dict()
      for key in model_state_dict.keys():
          assert key in loaded_model_state_dict
          model_value = model_state_dict[key]
          loaded_model_value = loaded_model_state_dict[key]
          assert (
              model_value.cpu().numpy()
              == loaded_model_value.cpu().numpy()
          ).all()
  
      optim_state_dict = optim.state_dict()
      for key in optim_state_dict.keys():
          assert key in loaded_optimizer_state_dict
          optim_value = optim_state_dict[key]
          loaded_optim_value = loaded_optimizer_state_dict[key]
          if idist.get_rank() == 0:
              assert optim_value == loaded_optim_value
  
      lr_scheduler_state_dict = lr_scheduler.state_dict()
      for key in lr_scheduler_state_dict.keys():
          assert key in loaded_lr_scheduler_state_dict
          lr_scheduler_value = lr_scheduler_state_dict[key]
          loaded_lr_scheduler_value = loaded_lr_scheduler_state_dict[
              key
          ]
          assert lr_scheduler_value == loaded_lr_scheduler_value
  
imports: import math
  import os
  from collections import defaultdict
  from unittest.mock import ANY, call, MagicMock, patch
  import clearml
  import pytest
  import torch
  from clearml.binding.frameworks import WeightsFileHandler
  from clearml.model import Framework
  import ignite.distributed
  from ignite.contrib.handlers.clearml_logger import ClearMLLogger, ClearMLSaver, global_step_from_engine, GradsHistHandler, GradsScalarHandler, OptimizerParamsHandler, OutputHandler, WeightsHistHandler, WeightsScalarHandler
  from ignite.engine import Engine, Events, State
  from ignite.handlers import Checkpoint
label_desc: Use image node_id 31 for calling a global function with example usage: _test_save_model_optimizer_lr_scheduler_with_state_dict(device, on_zero_rank) without return types