code: def count_accuracy(X, true_counts, air, batch_size):
      assert X.size(0) == true_counts.size(0), "Size mismatch."
      assert (
          X.size(0) % batch_size == 0
      ), "Input size must be multiple of batch_size."
      counts = torch.LongTensor(3, 4).zero_()
      error_latents = []
      error_indicators = []
  
      def count_vec_to_mat(vec, max_index):
          out = torch.LongTensor(vec.size(0), max_index + 1).zero_()
          out.scatter_(
              1, vec.type(torch.LongTensor).view(vec.size(0), 1), 1
          )
          return out
  
      for i in range(X.size(0) // batch_size):
          X_batch = X[i * batch_size : (i + 1) * batch_size]
          true_counts_batch = true_counts[
              i * batch_size : (i + 1) * batch_size
          ]
          z_where, z_pres = air.guide(X_batch, batch_size)
          inferred_counts = sum(z.cpu() for z in z_pres).squeeze().data
          true_counts_m = count_vec_to_mat(true_counts_batch, 2)
          inferred_counts_m = count_vec_to_mat(inferred_counts, 3)
          counts += torch.mm(true_counts_m.t(), inferred_counts_m)
          error_ind = 1 - (true_counts_batch == inferred_counts).long()
          error_ix = error_ind.nonzero(as_tuple=False).squeeze()
          error_latents.append(
              latents_to_tensor((z_where, z_pres)).index_select(
                  0, error_ix
              )
          )
          error_indicators.append(error_ind)
  
      acc = counts.diag().sum().float() / X.size(0)
      error_indices = (
          torch.cat(error_indicators).nonzero(as_tuple=False).squeeze()
      )
      if X.is_cuda:
          error_indices = error_indices.cuda()
      return acc, counts, torch.cat(error_latents), error_indices
  
imports: import argparse
  import math
  import os
  import time
  from functools import partial
  import numpy
  import torch
  import visdom
  from air import AIR, latents_to_tensor
  from viz import draw_many, tensor_to_objs
  import pyro
  import pyro.contrib.examples.multi_mnist
  import pyro.optim
  import pyro.poutine
  from pyro.contrib.examples.util import get_data_directory
  from pyro.infer import SVI, JitTraceGraph_ELBO, TraceGraph_ELBO
label_desc: Use image node_id 1 for calling a global function with example usage: count_accuracy(X, true_counts, air, batch_size) and returns: acc,  counts,  torch,  error_indices, out