diff --git a/requirements.txt b/requirements.txt
index cb1a66eaf..4d88f8fe7 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -12,4 +12,4 @@ OpenEXR >= 1.3.2
 termcolor >= 1.1.0
 trimesh >= 2.37.22
 # Required by trimesh.
-networkx
+networkx
\ No newline at end of file
diff --git a/tensorflow_graphics/util/test_case.py b/tensorflow_graphics/util/test_case.py
index f1ac35545..61ff39f3a 100644
--- a/tensorflow_graphics/util/test_case.py
+++ b/tensorflow_graphics/util/test_case.py
@@ -31,6 +31,10 @@
 import tensorflow as tf
 
 from tensorflow_graphics.util import tfg_flags
+# pylint: disable=g-direct-tensorflow-import
+from ai-edge-litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import
+
 
 FLAGS = flags.FLAGS
 
@@ -98,6 +102,7 @@ def _compute_gradient_error(self, x, y, x_init, delta=1e-6):
       error = 0
       row_max_error = 0
       column_max_error = 0
+      max_error = 0
       for j_t, j_n in grad:
         if j_t.size or j_n.size:  # Handle zero size tensors correctly
           diff = np.fabs(j_t - j_n)
@@ -364,7 +369,7 @@ def assert_tf_lite_convertible(self,
             sess, in_tensors, out_tensors)
         tflite_model = converter.convert()
         # Load TFLite model and allocate tensors.
-        interpreter = tf.lite.Interpreter(model_content=tflite_model)
+        interpreter = tfl_interpreter.Interpreter(model_content=tflite_model)
         interpreter.allocate_tensors()
         # If no test inputs provided then randomly generate inputs.
         if test_inputs is None: