diff --git a/.github/workflows/make_wheel_Linux_x86.sh b/.github/workflows/make_wheel_Linux_x86.sh
index 5a9a5ce9..69de8a04 100644
--- a/.github/workflows/make_wheel_Linux_x86.sh
+++ b/.github/workflows/make_wheel_Linux_x86.sh
@@ -56,6 +56,8 @@ if [[ "$TF_VERSION" =~ ^2\.1[3-9]\.[0-9]$ ]] ; then
   export PROTOBUF_VERSION='4.23.4'
 fi
 
+docker system prune -f
+
 DOCKER_BUILDKIT=1 docker build --no-cache \
     -f tools/docker/build_wheel.Dockerfile \
     --output type=local,dest=wheelhouse \
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_ops_test.py b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_ops_test.py
index 8723daa0..5f479252 100644
--- a/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_ops_test.py
+++ b/tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_ops_test.py
@@ -81,6 +81,52 @@
     from tensorflow.keras.legacy.optimizers import Adam
 
 
+class MyModel(layers.Layer):
+
+  def __init__(self):
+    super().__init__()
+    self.embeddings = de.get_variable(name="p1")
+
+  @tf.function
+  def call(self, indices_input, ids_input):
+    shape_input = tf.constant([2, 3, 4], dtype=tf.int64)
+
+    sparse_ids = tf.sparse.SparseTensor(indices=indices_input,
+                                        values=ids_input,
+                                        dense_shape=shape_input)
+
+    embeddings_result = de.safe_embedding_lookup_sparse(self.embeddings,
+                                                        sparse_ids,
+                                                        name="safe_sp_emb",
+                                                        return_trainable=False)
+    return embeddings_result
+
+
+# class MyModel(tf.Module):
+#   def __init__(self):
+#     super().__init__()
+#     # Create the variable as an attribute of the module
+#     self.embeddings = de.get_variable(name="p1")
+#
+#   @tf.function
+#   def __call__(self, indices_input, ids_input):
+#     shape_input = tf.constant([2, 3, 4], dtype=tf.int64)
+#
+#     sparse_ids = tf.sparse.SparseTensor(
+#       indices=indices_input,
+#       values=ids_input,
+#       dense_shape=shape_input
+#     )
+#
+#     embeddings_result = de.safe_embedding_lookup_sparse(
+#       self.embeddings,
+#       sparse_ids,
+#       name="safe_sp_emb",
+#       return_trainable=False
+#     )
+#     return embeddings_result
+
+
 # pylint: disable=missing-class-docstring
 # pylint: disable=missing-function-docstring
 def _type_converter(tf_type):
@@ -570,20 +616,6 @@ def test_scope_reuse_sparse_embedding_lookup(self):
     self.assertAllEqual(p2._tables[0].name, "test_p2_mht_1of1")
 
   def test_keras_input_safe_sparse_embedding_lookup(self):
-    embeddings = de.get_variable(name="p1")
-
-    @tf.function
-    def model(indices_input, ids_input):
-      shape_input = tf.constant([2, 3, 4], dtype=tf.int64)
-
-      sparse_ids = tf.sparse.SparseTensor(indices=indices_input,
-                                          values=ids_input,
-                                          dense_shape=shape_input)
-
-      embeddings_result = de.safe_embedding_lookup_sparse(
-          embeddings, sparse_ids, name="safe_sp_emb", return_trainable=False)
-      return embeddings_result
-
     indices_data = tf.constant([
         [0, 0, 0],
         [0, 0, 1],
@@ -593,9 +625,10 @@ def model(indices_input, ids_input):
         [1, 1, 0],
         [1, 1, 1],
     ],
+                               name="indices_data",
                                dtype=tf.int64)
-    ids_data = tf.constant([0, 1, -1, -1, 2, 0, 1], dtype=tf.int64)
-
+    ids_data = tf.constant([0, 1, -1, -1, 2, 0, 1], name="ids", dtype=tf.int64)
+    model = MyModel()
     # output = model(indices_data, ids_data)
 
   def test_scope_reuse_safe_sparse_embedding_lookup(self):
diff --git a/tools/testing/build_and_run_tests.sh b/tools/testing/build_and_run_tests.sh
index 59cdb74c..05920f10 100644
--- a/tools/testing/build_and_run_tests.sh
+++ b/tools/testing/build_and_run_tests.sh
@@ -31,6 +31,8 @@ if [ "$TF_VERSION" = "2.6.3" ] ; then
   python -m pip install numpy==1.19.5 --force-reinstall
 fi
 
+bazel clean --expunge
+
 # Avoid SystemError: initialization of _pywrap_checkpoint_reader raised unreported exception
 pip install tensorflow==$TF_VERSION
 
@@ -79,8 +81,9 @@ fi
 
 python -m pytest -v -s --functions-durations=20 --modules-durations=5 $IGNORE_HKV $SKIP_CUSTOM_OP_TESTS_FLAG $EXTRA_ARGS ./tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/
 
-
-
 # Release disk space
 bazel clean --expunge
 rm -f ./tensorflow_recommenders_addons/dynamic_embedding/core/_*_ops.so
+rm -rf /tmp/*
+apt-get clean && rm -rf /var/lib/apt/lists/*
+rm -rf /var/log/*
\ No newline at end of file