code: def test_serialization(jit, feature_dim, outcome_dist):
      if torch.__version__.startswith("2") and jit:
          pytest.xfail(
              reason="https://github.com/pyro-ppl/pyro/issues/3221"
          )
      x, t, y = generate_data(num_data=32, feature_dim=feature_dim)
      if outcome_dist == "exponential":
          y.clamp_(min=1e-20)
      cevae = CEVAE(
          feature_dim,
          outcome_dist=outcome_dist,
          num_samples=1000,
          hidden_dim=32,
      )
      cevae.fit(x, t, y, num_epochs=4, batch_size=8)
      pyro.set_rng_seed(0)
      expected_ite = cevae.ite(x)
  
      if jit:
          traced_cevae = cevae.to_script_module()
          f = io.BytesIO()
          torch.jit.save(traced_cevae, f)
          f.seek(0)
          loaded_cevae = torch.jit.load(f)
      else:
          f = io.BytesIO()
          with warnings.catch_warnings():
              warnings.filterwarnings("ignore", category=UserWarning)
              torch.save(cevae, f)
          f.seek(0)
          loaded_cevae = torch.load(f)
  
      pyro.set_rng_seed(0)
      actual_ite = loaded_cevae.ite(x)
      assert_close(actual_ite, expected_ite, atol=0.1)
  
file_path: pyro/tests/contrib/cevae/test_cevae.py
imports: import io
  import warnings
  import pytest
  import torch
  import pyro
  import pyro.distributions
  from pyro.contrib.cevae import CEVAE, DistributionNet
  from tests.common import assert_close
label_desc: Use image node_id 3 for calling a global function with example usage: test_serialization(jit, feature_dim, outcome_dist) without return types