code: def supervised_training_step_apex(
      model: torch.nn.Module,
      optimizer: torch.optim.Optimizer,
      loss_fn: Union[
          Callable[[Any, Any], torch.Tensor], torch.nn.Module
      ],
      device: Optional[Union[str, torch.device]] = None,
      non_blocking: bool = False,
      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(),
      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 apex.
  
      Args:
          model: the model to train.
          optimizer: the optimizer to use.
          loss_fn: the loss function that receives `y_pred` and `y`, and returns the loss as a tensor.
          device: device type specification (default: None).
              Applies to batches after starting the engine. Model *will not* be moved.
              Device can be CPU, GPU.
          non_blocking: if True and this copy is between CPU and GPU, the copy may occur asynchronously
              with respect to the host. For other cases, this argument has no effect.
          prepare_batch: function that receives `batch`, `device`, `non_blocking` and outputs
              tuple of tensors `(batch_x, batch_y)`.
          model_transform: function that receives the output from the model and convert it into the form as required
              by the loss function
          output_transform: function that receives 'x', 'y', 'y_pred', 'loss' and returns value
              to be assigned to engine's state.output after each iteration. Default is returning `loss.item()`.
          gradient_accumulation_steps: Number of steps the gradients should be accumulated across.
              (default: 1 (means no gradient accumulation))
          model_fn: the model function that receives `model` and `x`, and returns `y_pred`.
  
      Returns:
          Callable: update function.
  
      Examples:
          .. code-block:: python
  
              from ignite.engine import Engine, supervised_training_step_apex
  
              model = ...
              optimizer = ...
              loss_fn = ...
  
              update_fn = supervised_training_step_apex(model, optimizer, loss_fn, 'cuda')
              trainer = Engine(update_fn)
  
      .. versionadded:: 0.4.5
      .. versionchanged:: 0.4.7
          Added Gradient Accumulation.
      .. versionchanged:: 0.4.11
          Added `model_transform` to transform model's output
      .. versionchanged:: 0.4.13
          Added `model_fn` to customize model's application on the sample
      """
  
      try:
          from apex import amp as apex_amp
      except ModuleNotFoundError:
          raise ModuleNotFoundError(
              "Please install apex from https://github.com/nvidia/apex to use amp_mode='apex'."
          )
  
      if gradient_accumulation_steps <= 0:
          raise ValueError(
              "Gradient_accumulation_steps must be strictly positive. "
              "No gradient accumulation if the value set to one (default)."
          )
  
      def update(
          engine: Engine, batch: Sequence[torch.Tensor]
      ) -> Union[Any, Tuple[torch.Tensor]]:
          if (
              engine.state.iteration - 1
          ) % gradient_accumulation_steps == 0:
              optimizer.zero_grad()
          model.train()
          x, y = prepare_batch(
              batch, device=device, non_blocking=non_blocking
          )
          output = model_fn(model, x)
          y_pred = model_transform(output)
          loss = loss_fn(y_pred, y)
          if gradient_accumulation_steps > 1:
              loss = loss / gradient_accumulation_steps
          with apex_amp.scale_loss(loss, optimizer) as scaled_loss:
              scaled_loss.backward()
          if engine.state.iteration % gradient_accumulation_steps == 0:
              optimizer.step()
          return output_transform(
              x, y, y_pred, loss * gradient_accumulation_steps
          )
  
      return update
  
file_path: ignite/ignite/engine/__init__.py
imports: from collections.abc import Mapping
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union
  import torch
  import ignite.distributed
  from ignite.engine.deterministic import DeterministicEngine
  from ignite.engine.engine import Engine
  from ignite.engine.events import CallableEventWithFilter, EventEnum, Events, EventsList, RemovableEventHandle, State
  from ignite.metrics import Metric
  from ignite.utils import convert_tensor
label_desc: Use image node_id 4 for calling a global function with example usage: supervised_training_step_apex(model, optimizer, loss_fn, device, non_blocking, prepare_batch, model_transform, output_transform, gradient_accumulation_steps, model_fn) and returns: update, output_transform