code: def wrap_training_step(
      wrapped, instance: "pl.LightningModule", args, kwargs
  ):
      output_dict = wrapped(*args, **kwargs)
  
      if (
          isinstance(output_dict, dict)
          and output_dict is not None
          and "log" in output_dict
      ):
          log_dict = output_dict.pop("log")
          instance.log_dict(log_dict, on_step=True)
  
      return output_dict
  
imports: import copy
  import importlib
  import os
  from dataclasses import dataclass, is_dataclass
  from enum import Enum
  from functools import lru_cache
  from pathlib import Path
  from typing import List, Optional, Tuple, Union
  import wrapt
  from nemo.utils import AppState, logging
  from nemo.utils.data_utils import resolve_cache_dir
  from nemo.utils.data_utils import is_datastore_path
label_desc: Use image node_id 6 for calling a global function with example usage: wrap_training_step(wrapped, instance, args, kwargs) and returns: output_dict