code: def run_on_device(f, device):
      """Runs `f` under a tf.device context on the given device."""
      f = utils.smart_autograph(f)
  
      @tf.autograph.experimental.do_not_convert
      @functools.wraps(f)
      def wrapper(*args, **kwargs):
          with tf.device(device):
              args = tree.map_structure(tf.identity, args)
              kwargs = tree.map_structure(tf.identity, kwargs)
              return f(*args, **kwargs)
  
      return wrapper
  
file_path: sonnet/sonnet/src/functional/utils.py
imports: import functools
  from sonnet.src import utils
  import tensorflow
  import tree
label_desc: Use image node_id 2 for calling a global function with example usage: run_on_device(f, device) and returns: wrapper, f