code: def test_model(model, test_dataset):
      correct = []
      for batch, (images, labels) in enumerate(test_dataset):
          logits = model(images)
          correct += list(
              labels.numpy() == np.argmax(logits.numpy(), axis=1)
          )
      return np.mean(correct)
  
file_path: auptimizer/Examples/compression/mnist_tensorflow/mnist.py
imports: from aup import print_result, aup_args
  import tensorflow
  import numpy
  from tensorflow import keras
  from tensorflow.keras import layers
  import tqdm
  from math import log
  import sys
  import aup
label_desc: Use image node_id 2 for calling a global function with example usage: test_model(model, test_dataset) and returns: np