code: def Model(args, data):
      """Dispatch between different model classes."""
      if args.heterogeneous:
          assert args.incubation_time == 0
          assert args.overdispersion == 0
          return HeterogeneousSIRModel(
              args.population, args.recovery_time, data
          )
      elif args.incubation_time > 0:
          assert args.incubation_time > 1
          if args.concentration < math.inf:
              return SuperspreadingSEIRModel(
                  args.population,
                  args.incubation_time,
                  args.recovery_time,
                  data,
              )
          elif args.overdispersion > 0:
              return OverdispersedSEIRModel(
                  args.population,
                  args.incubation_time,
                  args.recovery_time,
                  data,
              )
          else:
              return SimpleSEIRModel(
                  args.population,
                  args.incubation_time,
                  args.recovery_time,
                  data,
              )
      else:
          if args.concentration < math.inf:
              return SuperspreadingSIRModel(
                  args.population, args.recovery_time, data
              )
          elif args.overdispersion > 0:
              return OverdispersedSIRModel(
                  args.population, args.recovery_time, data
              )
          else:
              return SimpleSIRModel(
                  args.population, args.recovery_time, data
              )
  
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 1 for calling a global function with example usage: Model(args, data) and returns: HeterogeneousSIRModel, SuperspreadingSEIRModel, SuperspreadingSIRModel, OverdispersedSEIRModel, SimpleSEIRModel, OverdispersedSIRModel, SimpleSIRModel