diff --git a/keras_rs/src/layers/embedding/distributed_embedding_test.py b/keras_rs/src/layers/embedding/distributed_embedding_test.py
index 94967c1..c2149e2 100644
--- a/keras_rs/src/layers/embedding/distributed_embedding_test.py
+++ b/keras_rs/src/layers/embedding/distributed_embedding_test.py
@@ -598,8 +598,8 @@ def test_correctness(
                         non_trainable_layouts,
                     ),
                 )(
-                    layer.trainable_variables,
-                    layer.non_trainable_variables,
+                    [v.value for v in layer.trainable_variables],
+                    [v.value for v in layer.non_trainable_variables],
                     preprocessed,
                 )
             else: