code: def main(args):
      if args.cuda:
          torch.set_default_device("cuda")
  
      # Generate synthetic data.
      pyro.set_rng_seed(args.seed)
      x_train, t_train, y_train, _ = generate_data(args)
  
      # Train.
      pyro.set_rng_seed(args.seed)
      pyro.clear_param_store()
      cevae = CEVAE(
          feature_dim=args.feature_dim,
          latent_dim=args.latent_dim,
          hidden_dim=args.hidden_dim,
          num_layers=args.num_layers,
          num_samples=10,
      )
      cevae.fit(
          x_train,
          t_train,
          y_train,
          num_epochs=args.num_epochs,
          batch_size=args.batch_size,
          learning_rate=args.learning_rate,
          learning_rate_decay=args.learning_rate_decay,
          weight_decay=args.weight_decay,
      )
  
      # Evaluate.
      x_test, t_test, y_test, true_ite = generate_data(args)
      true_ate = true_ite.mean()
      print("true ATE = {:0.3g}".format(true_ate.item()))
      naive_ate = (
          y_test[t_test == 1].mean() - y_test[t_test == 0].mean()
      )
      print("naive ATE = {:0.3g}".format(naive_ate))
      if args.jit:
          cevae = cevae.to_script_module()
      est_ite = cevae.ite(x_test)
      est_ate = est_ite.mean()
      print("estimated ATE = {:0.3g}".format(est_ate.item()))
  
file_path: pyro/examples/contrib/cevae/synthetic.py
imports: import argparse
  import logging
  import torch
  import pyro
  import pyro.distributions
  from pyro.contrib.cevae import CEVAE
label_desc: Use image node_id 2 for calling a global function with example usage: main(args) without return types