code: def run_infer(
      original_model, sharded_model, data_gen_fn, output_transform_fn
  ):
      # prepare input
      data = data_gen_fn()
      data = {k: v.cuda() for k, v in data.items()}
      # run forward
      org_output = original_model(**data)
      org_output = output_transform_fn(org_output)
  
      shard_output = sharded_model(**data)
      shard_output = output_transform_fn(shard_output)
  
      return org_output, shard_output
  
label_desc: Use image node_id 2 for calling a global function with example usage: run_infer(original_model, sharded_model, data_gen_fn, output_transform_fn) and returns: org_output,  shard_output