code: def infer_mcmc(args, model):
      parallel = args.num_chains > 1
      energies = []
  
      def hook_fn(kernel, *unused):
          e = float(kernel._potential_energy_last)
          energies.append(e)
          if args.verbose:
              logging.info("potential = {:0.6g}".format(e))
  
      mcmc = model.fit_mcmc(
          heuristic_num_particles=args.smc_particles,
          heuristic_ess_threshold=args.ess_threshold,
          warmup_steps=args.warmup_steps,
          num_samples=args.num_samples,
          num_chains=args.num_chains,
          mp_context="spawn" if parallel else None,
          max_tree_depth=args.max_tree_depth,
          arrowhead_mass=args.arrowhead_mass,
          num_quant_bins=args.num_bins,
          haar=args.haar,
          haar_full_mass=args.haar_full_mass,
          jit_compile=args.jit,
          hook_fn=None if parallel else hook_fn,
      )
  
      mcmc.summary()
      if args.plot and energies:
          import matplotlib.pyplot as plt
  
          plt.figure(figsize=(6, 3))
          plt.plot(energies)
          plt.xlabel("MCMC step")
          plt.ylabel("potential energy")
          plt.title("MCMC energy trace")
          plt.tight_layout()
  
      return model.samples
  
file_path: pyro/examples/contrib/epidemiology/sir.py
imports: import argparse
  import logging
  import math
  import torch
  from torch.distributions import biject_to, constraints
  import pyro
  from pyro.contrib.epidemiology.models import HeterogeneousSIRModel, OverdispersedSEIRModel, OverdispersedSIRModel, SimpleSEIRModel, SimpleSIRModel, SuperspreadingSEIRModel, SuperspreadingSIRModel
label_desc: Use image node_id 3 for calling a global function with example usage: infer_mcmc(args, model) and returns: model