code: def batch_second(args, kwargs):
      def set_batch_second(ew):
          ew.set_batch_first(False)
  
      def reset_batch_first(ew):
          ew.set_batch_first(True)
  
      tree_map_only(ExpandedWeight, set_batch_second, args)
      tree_map_only(ExpandedWeight, set_batch_second, kwargs)
      try:
          yield
      finally:
          tree_map_only(ExpandedWeight, reset_batch_first, args)
          tree_map_only(ExpandedWeight, reset_batch_first, kwargs)
  
imports: from contextlib import contextmanager
  import torch
  import functools
  from torch._decomp import decomposition_table
  from typing import Callable, Dict
  from torch.utils._pytree import tree_map_only
label_desc: Use image node_id 1 for calling a global function with example usage: batch_second(args, kwargs) without return types