code: def one_hot_model(pseudocounts, classes=None):
      probs_prior = dist.Dirichlet(pseudocounts)
      probs = pyro.sample("probs", probs_prior)
      with pyro.plate(
          "classes",
          classes.size(0) if classes is not None else 1,
          dim=-1,
      ):
          return pyro.sample(
              "obs", dist.OneHotCategorical(probs), obs=classes
          )
  
file_path: pyro/tests/infer/test_predictive.py
imports: import pytest
  import torch
  import pyro
  import pyro.distributions
  import pyro.optim
  import pyro.poutine
  from pyro.infer import SVI, Predictive, Trace_ELBO
  from pyro.infer.autoguide import AutoDelta, AutoDiagonalNormal
  from tests.common import assert_close
label_desc: Use image node_id 2 for calling a global function with example usage: one_hot_model(pseudocounts, classes) and returns: pyro