code: def setup():
      global model
  
      class DoesUseSame(objax.Module):
          def __init__(self):
              fn = functools.partial(
                  wide_resnet.WideResNet, depth=28, width=6
              )
              self.model = fn(3 * 4, 2)
  
              model_vars = self.model.vars()
              self.ema = objax.optimizer.ExponentialMovingAverage(
                  model_vars, momentum=0.999, debias=True
              )
  
              def predict_op(x, y):
                  # The model takes SEVERAL images and checks if they all correspond
                  # to the same original image.
                  # Guaranteed that the first N-1 all do, the test is if the last does.
                  xx = jn.concatenate([jn.abs(x), jn.abs(y)], axis=1)
                  return self.model(xx, training=False)
  
              self.predict = objax.Jit(
                  self.ema.replace_vars(predict_op),
                  model_vars + self.ema.vars(),
              )
  
      model = DoesUseSame()
      checkpoint = objax.io.Checkpoint(
          "models/step2/", keep_ckpts=5, makedir=True
      )
      start_epoch, last_ckpt = checkpoint.restore(model.vars())
  
file_path: privacy/research/instahide_attack_2020/step_3_second_graph.py
imports: import objax
  import numpy
  import jax.numpy
  import functools
  import os
  import random
  from objax.zoo import wide_resnet
label_desc: Use image node_id 1 for calling a global function with example usage: setup() and returns: self