From a5209f888a872febb3cb4eb82cda85f63d7782b1 Mon Sep 17 00:00:00 2001
From: tfx-team <tensorflow-extended-nonhuman@googlegroups.com>
Date: Tue, 13 Aug 2024 13:22:17 -0700
Subject: [PATCH] change MessageFactory class to GetMessageClass()

PiperOrigin-RevId: 662629553
---
 tfx/utils/proto_utils.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tfx/utils/proto_utils.py b/tfx/utils/proto_utils.py
index de5abf4fd7..381b557274 100644
--- a/tfx/utils/proto_utils.py
+++ b/tfx/utils/proto_utils.py
@@ -95,9 +95,8 @@ def _create_proto_instance_from_name(
     message_name: str, pool: descriptor_pool.DescriptorPool) -> ProtoMessage:
   """Creates a protobuf message instance from a given message name."""
   message_descriptor = pool.FindMessageTypeByName(message_name)
-  factory = message_factory.MessageFactory(pool)
-  message_type = factory.GetPrototype(message_descriptor)
-  return message_type()
+  message_class = message_factory.GetMessageClass(message_descriptor)
+  return message_class()
 
 
 def get_pool_with_descriptors(