code: def main(args):
      pyro.set_rng_seed(0)
      pyro.clear_param_store()
  
      # load data
      if args.dataset == "dipper":
          capture_history_file = (
              os.path.dirname(os.path.abspath(__file__))
              + "/dipper_capture_history.csv"
          )
      elif args.dataset == "vole":
          capture_history_file = (
              os.path.dirname(os.path.abspath(__file__))
              + "/meadow_voles_capture_history.csv"
          )
      else:
          raise ValueError(
              "Available datasets are 'dipper' and 'vole'."
          )
  
      capture_history = torch.tensor(
          np.genfromtxt(capture_history_file, delimiter=",")
      ).float()[:, 1:]
      N, T = capture_history.shape
      print(
          "Loaded {} capture history for {} individuals collected over {} time periods.".format(
              args.dataset, N, T
          )
      )
  
      if args.dataset == "dipper" and args.model in ["4", "5"]:
          sex_file = (
              os.path.dirname(os.path.abspath(__file__))
              + "/dipper_sex.csv"
          )
          sex = torch.tensor(
              np.genfromtxt(sex_file, delimiter=",")
          ).float()[:, 1]
          print("Loaded dipper sex data.")
      elif args.dataset == "vole" and args.model in ["4", "5"]:
          raise ValueError(
              "Cannot run model_{} on meadow voles data, since we lack sex "
              "information for these animals.".format(args.model)
          )
      else:
          sex = None
  
      model = models[args.model]
  
      # we use poutine.block to only expose the continuous latent variables
      # in the models to AutoDiagonalNormal (all of which begin with 'phi'
      # or 'rho')
      def expose_fn(msg):
          return msg["name"][0:3] in ["phi", "rho"]
  
      # we use a mean field diagonal normal variational distributions (i.e. guide)
      # for the continuous latent variables.
      guide = AutoDiagonalNormal(
          poutine.block(model, expose_fn=expose_fn)
      )
  
      # since we enumerate the discrete random variables,
      # we need to use TraceEnum_ELBO or TraceTMC_ELBO.
      optim = Adam({"lr": args.learning_rate})
      if args.tmc:
          elbo = TraceTMC_ELBO(max_plate_nesting=1)
          tmc_model = poutine.infer_config(
              model,
              lambda msg: {
                  "num_samples": args.tmc_num_samples,
                  "expand": False,
              }
              if msg["infer"].get("enumerate", None) == "parallel"
              else {},
          )  # noqa: E501
          svi = SVI(tmc_model, guide, optim, elbo)
      else:
          elbo = TraceEnum_ELBO(
              max_plate_nesting=1,
              num_particles=20,
              vectorize_particles=True,
          )
          svi = SVI(model, guide, optim, elbo)
  
      losses = []
  
      print(
          "Beginning training of model_{} with Stochastic Variational Inference.".format(
              args.model
          )
      )
  
      for step in range(args.num_steps):
          loss = svi.step(capture_history, sex)
          losses.append(loss)
          if step % 20 == 0 and step > 0 or step == args.num_steps - 1:
              print(
                  "[iteration %03d] loss: %.3f"
                  % (step, np.mean(losses[-20:]))
              )
  
      # evaluate final trained model
      elbo_eval = TraceEnum_ELBO(
          max_plate_nesting=1,
          num_particles=2000,
          vectorize_particles=True,
      )
      svi_eval = SVI(model, guide, optim, elbo_eval)
      print(
          "Final loss: %.4f"
          % svi_eval.evaluate_loss(capture_history, sex)
      )
  
file_path: pyro/examples/capture_recapture/cjs.py
imports: import argparse
  import os
  import numpy
  import torch
  import pyro
  import pyro.distributions
  import pyro.poutine
  from pyro.infer import SVI, TraceEnum_ELBO, TraceTMC_ELBO
  from pyro.infer.autoguide import AutoDiagonalNormal
  from pyro.optim import Adam
label_desc: Use image node_id 6 for calling a global function with example usage: main(args) and returns: unknown