Class CustomModule

e2e/integration_tests/convert_predict.py:313–332

    save_dir: directory name of where the saved model will be stored.
  """
  class CustomModule(tf.Module):

    def __init__(self):
        super(CustomModule, self).__init__()

    @tf.function(input_signature=[tf.TensorSpec([], tf.int32),
                                  tf.TensorSpec([], tf.int32),
                                  tf.TensorSpec([], tf.int32)])
    def control_flow(self, x, y, z):
        i = 0
        while i < z:
            i += 1
            j = 0
            while j < y:
                j += 1
                if z > 0:
                    x += 1
                else:
                    x += 2
        return x


  module = CustomModule()