Helper function that computes the current class prediction
  :param samples: numpy array with input samples (dims must match x)
  :return: the argmax output of predictions, i.e. the current predicted class

def model_argmax(model, samples):
  """
  Helper function that computes the current class prediction
  :param samples: numpy array with input samples (dims must match x)
  :return: the argmax output of predictions, i.e. the current predicted class
  """
  tfe = tf.contrib.eager
  tf_samples = tfe.Variable(samples)
  probabilities = model.get_probs(tf_samples)

  if samples.shape[0] == 1:
    return tf.argmax(probabilities)
  else:
    return tf.argmax(probabilities, axis=1)