diff --git a/compiler/tests/python/test_round_trip.py b/compiler/tests/python/test_round_trip.py
index aab3d9f75..4f35c48e9 100644
--- a/compiler/tests/python/test_round_trip.py
+++ b/compiler/tests/python/test_round_trip.py
@@ -11,8 +11,8 @@ VALID_INPUTS = [
     }
     """,
     """
-    func @apply_lookup_table(%arg0: !HLFHE.eint<2>, %arg1: memref<4xi2>) -> !HLFHE.eint<2> {
-        %1 = "HLFHE.apply_lookup_table"(%arg0, %arg1): (!HLFHE.eint<2>, memref<4xi2>) -> (!HLFHE.eint<2>)
+    func @apply_lookup_table(%arg0: !HLFHE.eint<2>, %arg1: tensor<4xi2>) -> !HLFHE.eint<2> {
+        %1 = "HLFHE.apply_lookup_table"(%arg0, %arg1): (!HLFHE.eint<2>, tensor<4xi2>) -> (!HLFHE.eint<2>)
         return %1: !HLFHE.eint<2>
     }
     """,