code: def main(args):
      funsor.set_backend("torch")
  
      # XXX Temporary fix after https://github.com/pyro-ppl/pyro/pull/2701
      import pyro
  
      pyro.enable_validation(False)
  
      encoder = Encoder()
      decoder = Decoder()
  
      # These rely on type hints on the .forward() methods.
      encode = funsor.function(encoder)
      decode = funsor.function(decoder)
  
      @funsor.montecarlo.MonteCarlo()
      def loss_function(data, subsample_scale):
          # Lazily sample from the guide.
          loc, scale = encode(data)
          q = funsor.Independent(
              dist.Normal(loc["i"], scale["i"], value="z_i"),
              "z",
              "i",
              "z_i",
          )
  
          # Evaluate the model likelihood at the lazy value z.
          probs = decode("z")
          p = dist.Bernoulli(probs["x", "y"], value=data["x", "y"])
          p = p.reduce(ops.add, {"x", "y"})
  
          # Construct an elbo. This is where sampling happens.
          elbo = funsor.Integrate(q, p - q, "z")
          elbo = elbo.reduce(ops.add, "batch") * subsample_scale
          loss = -elbo
          return loss
  
      train_loader = torch.utils.data.DataLoader(
          MNIST(
              DATA_PATH,
              train=True,
              download=True,
              transform=transforms.ToTensor(),
          ),
          batch_size=args.batch_size,
          shuffle=True,
      )
  
      encoder.train()
      decoder.train()
      optimizer = optim.Adam(
          list(encoder.parameters()) + list(decoder.parameters()),
          lr=1e-3,
      )
      for epoch in range(args.num_epochs):
          train_loss = 0
          for batch_idx, (data, _) in enumerate(train_loader):
              subsample_scale = float(
                  len(train_loader.dataset) / len(data)
              )
              data = data[:, 0, :, :]
              data = funsor.Tensor(
                  data, OrderedDict(batch=Bint[len(data)])
              )
  
              optimizer.zero_grad()
              loss = loss_function(data, subsample_scale)
              assert isinstance(loss, funsor.Tensor), loss.pretty()
              loss.data.backward()
              train_loss += loss.item()
              optimizer.step()
              if batch_idx % 50 == 0:
                  print(f"  loss = {loss.item()}")
                  if batch_idx and args.smoke_test:
                      return
          print(f"epoch {epoch} train_loss = {train_loss}")
  
imports: import argparse
  import os
  import typing
  from collections import OrderedDict
  import torch
  import torch.utils.data
  from torch import nn, optim
  from torch.nn import functional
  from torchvision import transforms
  from torchvision.datasets import MNIST
  import funsor
  import funsor.ops
  import funsor.torch.distributions
  from funsor.domains import Bint, Reals
label_desc: Use image node_id 1 for calling a global function with example usage: main(args) and returns: loss