diff --git a/bindings/python/tests/testdata/generate_tflite.py b/bindings/python/tests/testdata/generate_tflite.py
index 61fc87b..a325089 100644
--- a/bindings/python/tests/testdata/generate_tflite.py
+++ b/bindings/python/tests/testdata/generate_tflite.py
@@ -20,7 +20,8 @@
 concrete_func = model.__call__.get_concrete_function(
     tf.TensorSpec(shape=[4, 3], dtype=tf.float32))
 
-converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
+converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func],
+                                                            model)
 tflite_model = converter.convert()
 
 this_dir = os.path.dirname(__file__)
diff --git a/integrations/tensorflow/bindings/python/iree/tf/support/module_utils.py b/integrations/tensorflow/bindings/python/iree/tf/support/module_utils.py
index 2f1e461..bd13b74 100644
--- a/integrations/tensorflow/bindings/python/iree/tf/support/module_utils.py
+++ b/integrations/tensorflow/bindings/python/iree/tf/support/module_utils.py
@@ -620,7 +620,8 @@
   for method, method_name in zip(methods, method_names):
     logging.info("Attempting to convert '%s' to tflite...", method_name)
     try:
-      converter = tf.lite.TFLiteConverter.from_concrete_functions([method])
+      converter = tf.lite.TFLiteConverter.from_concrete_functions([method],
+                                                                  module_class)
       logging.info("...converted '%s' to tflite.", method_name)
       tflite_modules.append(converter.convert())
     except Exception as e: