38e9391dcb5e8991a1a663601ce78bf7f08df474,tensorflow_transform/pretrained_models_test.py,PretrainedModelsTest,testApplyFunctionWithCheckpointTwoInputs,#PretrainedModelsTest#,155

Before Change

  def testApplyFunctionWithCheckpointTwoInputs(self):
    checkpoint = os.path.join(self.get_temp_dir(), "checkpoint_two")
    self.save_checkpoint_with_two_inputs(checkpoint)
    with self.test_session(graph=tf.Graph()) as sess:
      input1 = tf.compat.v1.placeholder(
          dtype=tf.int64, shape=[3], name="input1")
      input2 = tf.compat.v1.placeholder(
          dtype=tf.int64, shape=[3], name="input2")
      output_tensor = pretrained_models.apply_function_with_checkpoint(
          self.make_tensor_fn_two_inputs(), [input1, input2], checkpoint)
      feed_dict = {input1: [1, 2, 3], input2: [3, 2, 1]}
      output_value = sess.run(output_tensor, feed_dict=feed_dict)
      // [1, 2, 3] + [1, 2, 3] - [3, 2, 1] = [-1, 2, 5]
      self.assertAllEqual(output_value, [-1, 2, 5])


if __name__ == "__main__":
  // TODO(b/133440043): Remove this once TFT supports eager execution.
  tf.compat.v1.disable_eager_execution()
  tf.test.main()

After Change

  def testApplyFunctionWithCheckpointTwoInputs(self):
    checkpoint = os.path.join(self.get_temp_dir(), "checkpoint_two")
    self.save_checkpoint_with_two_inputs(checkpoint)
    with tf.compat.v1.Graph().as_default() as graph:
      with self.test_session(graph=graph) as sess:
        input1 = tf.compat.v1.placeholder(
            dtype=tf.int64, shape=[3], name="input1")
        input2 = tf.compat.v1.placeholder(
            dtype=tf.int64, shape=[3], name="input2")
        output_tensor = pretrained_models.apply_function_with_checkpoint(
            self.make_tensor_fn_two_inputs(), [input1, input2], checkpoint)
        feed_dict = {input1: [1, 2, 3], input2: [3, 2, 1]}
        output_value = sess.run(output_tensor, feed_dict=feed_dict)
        // [1, 2, 3] + [1, 2, 3] - [3, 2, 1] = [-1, 2, 5]
        self.assertAllEqual(output_value, [-1, 2, 5])


if __name__ == "__main__":
  tf.test.main()

Instances