import tensorflow as tf

      model = revnet.RevNet(config=config)
      if defun:
        # TODO(apassos): reenable after cond lets you return None
        model.call = tfe.defun(model.call)
      batch_size = 64
      num_burn = 5
      num_iters = 10
      with tf.device(device):
        images, _ = random_batch(batch_size, config)
        for _ in range(num_burn):
          model(images, training=False)
        if execution_mode:
          tfe.async_wait()
        gc.collect()
