diff --git a/.github/workflows/cpu-horovod.yml b/.github/workflows/cpu-horovod.yml
index 4cf4dadd01..481ee11163 100644
--- a/.github/workflows/cpu-horovod.yml
+++ b/.github/workflows/cpu-horovod.yml
@@ -43,7 +43,7 @@ jobs:
           branch=main
           if [[ $ref_type == "tag"* ]]
           then
-            raw=$(git branch -r --contains ${{ github.ref_name }})
-            branch=${raw/origin\/}
+            git -c protocol.version=2 fetch --no-tags --prune --progress --no-recurse-submodules --depth=1 origin +refs/heads/release*:refs/remotes/origin/release*
+            branch=$(git branch -r --contains ${{ github.ref_name }} --list '*release*' --format "%(refname:short)" | sed -e 's/^origin\///')
           fi
           cd ${{ github.workspace }}; tox -e py38-horovod-cpu -- $branch
diff --git a/.github/workflows/cpu-nvtabular.yml b/.github/workflows/cpu-nvtabular.yml
index c007ad5270..1d853ee4e2 100644
--- a/.github/workflows/cpu-nvtabular.yml
+++ b/.github/workflows/cpu-nvtabular.yml
@@ -38,7 +38,7 @@ jobs:
         branch=main
         if [[ $ref_type == "tag"* ]]
         then
-          raw=$(git branch -r --contains ${{ github.ref_name }})
-          branch=${raw/origin\/}
+          git -c protocol.version=2 fetch --no-tags --prune --progress --no-recurse-submodules --depth=1 origin +refs/heads/release*:refs/remotes/origin/release*
+          branch=$(git branch -r --contains ${{ github.ref_name }} --list '*release*' --format "%(refname:short)" | sed -e 's/^origin\///')
         fi
         GIT_COMMIT=`git rev-parse HEAD` tox -e py38-nvtabular-cpu -- $branch
diff --git a/.github/workflows/cpu-systems.yml b/.github/workflows/cpu-systems.yml
index 8e31a65c0f..7cdc9cb25c 100644
--- a/.github/workflows/cpu-systems.yml
+++ b/.github/workflows/cpu-systems.yml
@@ -38,7 +38,7 @@ jobs:
         branch=main
         if [[ $ref_type == "tag"* ]]
         then
-          raw=$(git branch -r --contains ${{ github.ref_name }})
-          branch=${raw/origin\/}
+          git -c protocol.version=2 fetch --no-tags --prune --progress --no-recurse-submodules --depth=1 origin +refs/heads/release*:refs/remotes/origin/release*
+          branch=$(git branch -r --contains ${{ github.ref_name }} --list '*release*' --format "%(refname:short)" | sed -e 's/^origin\///')
         fi
         GIT_COMMIT=`git rev-parse HEAD` tox -e py38-systems-cpu -- $branch
diff --git a/.github/workflows/datasets.yml b/.github/workflows/datasets.yml
index ae2c240fa7..7c0ac3c01c 100644
--- a/.github/workflows/datasets.yml
+++ b/.github/workflows/datasets.yml
@@ -32,8 +32,8 @@ jobs:
           branch=main
           if [[ $ref_type == "tag"* ]]
           then
-            raw=$(git branch -r --contains ${{ github.ref_name }})
-            branch=${raw/origin\/}
+            git -c protocol.version=2 fetch --no-tags --prune --progress --no-recurse-submodules --depth=1 origin +refs/heads/release*:refs/remotes/origin/release*
+            branch=$(git branch -r --contains ${{ github.ref_name }} --list '*release*' --format "%(refname:short)" | sed -e 's/^origin\///')
           fi
           pip install "pandas>=1.2.0,<1.4.0dev0"
           pip install "NVTabular@git+https://github.com/NVIDIA-Merlin/NVTabular.git@$branch"
diff --git a/.github/workflows/gpu-ci.yml b/.github/workflows/gpu-ci.yml
index d832a1bdb5..0b3ec20862 100644
--- a/.github/workflows/gpu-ci.yml
+++ b/.github/workflows/gpu-ci.yml
@@ -24,7 +24,7 @@ jobs:
         branch=main
         if [[ $ref_type == "tag"* ]]
         then
-          raw=$(git branch -r --contains ${{ github.ref_name }})
-          branch=${raw/origin\/}
+          git -c protocol.version=2 fetch --no-tags --prune --progress --no-recurse-submodules --depth=1 origin +refs/heads/release*:refs/remotes/origin/release*
+          branch=$(git branch -r --contains ${{ github.ref_name }} --list '*release*' --format "%(refname:short)" | sed -e 's/^origin\///')
         fi
         cd ${{ github.workspace }}; tox -e py38-gpu -- $branch
diff --git a/.github/workflows/implicit.yml b/.github/workflows/implicit.yml
index 42f751fcd9..3feac70d33 100644
--- a/.github/workflows/implicit.yml
+++ b/.github/workflows/implicit.yml
@@ -32,8 +32,8 @@ jobs:
           branch=main
           if [[ $ref_type == "tag"* ]]
           then
-            raw=$(git branch -r --contains ${{ github.ref_name }})
-            branch=${raw/origin\/}
+            git -c protocol.version=2 fetch --no-tags --prune --progress --no-recurse-submodules --depth=1 origin +refs/heads/release*:refs/remotes/origin/release*
+            branch=$(git branch -r --contains ${{ github.ref_name }} --list '*release*' --format "%(refname:short)" | sed -e 's/^origin\///')
           fi
           pip install "pandas>=1.2.0,<1.4.0dev0"
           pip install "NVTabular@git+https://github.com/NVIDIA-Merlin/NVTabular.git@$branch"
diff --git a/.github/workflows/lightfm.yml b/.github/workflows/lightfm.yml
index 715e64aaaa..05328c18d2 100644
--- a/.github/workflows/lightfm.yml
+++ b/.github/workflows/lightfm.yml
@@ -32,8 +32,8 @@ jobs:
           branch=main
           if [[ $ref_type == "tag"* ]]
           then
-            raw=$(git branch -r --contains ${{ github.ref_name }})
-            branch=${raw/origin\/}
+            git -c protocol.version=2 fetch --no-tags --prune --progress --no-recurse-submodules --depth=1 origin +refs/heads/release*:refs/remotes/origin/release*
+            branch=$(git branch -r --contains ${{ github.ref_name }} --list '*release*' --format "%(refname:short)" | sed -e 's/^origin\///')
           fi
           pip install "pandas>=1.2.0,<1.4.0dev0"
           pip install "NVTabular@git+https://github.com/NVIDIA-Merlin/NVTabular.git@$branch"
diff --git a/.github/workflows/pytorch.yml b/.github/workflows/pytorch.yml
index 0b341f2e9b..647680da0a 100644
--- a/.github/workflows/pytorch.yml
+++ b/.github/workflows/pytorch.yml
@@ -1,3 +1,4 @@
+
 name: pytorch
 
 on:
@@ -32,8 +33,8 @@ jobs:
           branch=main
           if [[ $ref_type == "tag"* ]]
           then
-            raw=$(git branch -r --contains ${{ github.ref_name }})
-            branch=${raw/origin\/}
+            git -c protocol.version=2 fetch --no-tags --prune --progress --no-recurse-submodules --depth=1 origin +refs/heads/release*:refs/remotes/origin/release*
+            branch=$(git branch -r --contains ${{ github.ref_name }} --list '*release*' --format "%(refname:short)" | sed -e 's/^origin\///')
           fi
           pip install "pandas>=1.2.0,<1.4.0dev0"
           pip install "NVTabular@git+https://github.com/NVIDIA-Merlin/NVTabular.git@$branch"
diff --git a/.github/workflows/tensorflow.yml b/.github/workflows/tensorflow.yml
index e58918e099..42cf1ab4de 100644
--- a/.github/workflows/tensorflow.yml
+++ b/.github/workflows/tensorflow.yml
@@ -38,8 +38,8 @@ jobs:
           branch=main
           if [[ $ref_type == "tag"* ]]
           then
-            raw=$(git branch -r --contains ${{ github.ref_name }})
-            branch=${raw/origin\/}
+            git -c protocol.version=2 fetch --no-tags --prune --progress --no-recurse-submodules --depth=1 origin +refs/heads/release*:refs/remotes/origin/release*
+            branch=$(git branch -r --contains ${{ github.ref_name }} --list '*release*' --format "%(refname:short)" | sed -e 's/^origin\///')
           fi
           pip install "pandas>=1.2.0,<1.4.0dev0"
           pip install "NVTabular@git+https://github.com/NVIDIA-Merlin/NVTabular.git@$branch"
@@ -93,8 +93,8 @@ jobs:
           branch=main
           if [[ $ref_type == "tag"* ]]
           then
-            raw=$(git branch -r --contains ${{ github.ref_name }})
-            branch=${raw/origin\/}
+            git -c protocol.version=2 fetch --no-tags --prune --progress --no-recurse-submodules --depth=1 origin +refs/heads/release*:refs/remotes/origin/release*
+            branch=$(git branch -r --contains ${{ github.ref_name }} --list '*release*' --format "%(refname:short)" | sed -e 's/^origin\///')
           fi
           pip install "pandas>=1.2.0,<1.4.0dev0"
           pip install "NVTabular@git+https://github.com/NVIDIA-Merlin/NVTabular.git@$branch"
diff --git a/.github/workflows/xgboost.yml b/.github/workflows/xgboost.yml
index 8a0b2b2003..fc92c618b3 100644
--- a/.github/workflows/xgboost.yml
+++ b/.github/workflows/xgboost.yml
@@ -32,8 +32,8 @@ jobs:
           branch=main
           if [[ $ref_type == "tag"* ]]
           then
-            raw=$(git branch -r --contains ${{ github.ref_name }})
-            branch=${raw/origin\/}
+            git -c protocol.version=2 fetch --no-tags --prune --progress --no-recurse-submodules --depth=1 origin +refs/heads/release*:refs/remotes/origin/release*
+            branch=$(git branch -r --contains ${{ github.ref_name }} --list '*release*' --format "%(refname:short)" | sed -e 's/^origin\///')
           fi
           pip install "pandas>=1.2.0,<1.4.0dev0"
           pip install "NVTabular@git+https://github.com/NVIDIA-Merlin/NVTabular.git@$branch"
diff --git a/merlin/models/tf/inputs/embedding.py b/merlin/models/tf/inputs/embedding.py
index b253c48004..f1c49e792a 100644
--- a/merlin/models/tf/inputs/embedding.py
+++ b/merlin/models/tf/inputs/embedding.py
@@ -618,6 +618,7 @@ def _get_dim(col, embedding_dims, infer_dim_fn):
     return dim
 
 
+@tf.keras.utils.register_keras_serializable(package="merlin.models")
 class AverageEmbeddingsByWeightFeature(tf.keras.layers.Layer):
     def __init__(self, weight_feature_name: str, axis=1, **kwargs):
         """Computes the weighted average of a Tensor based
@@ -694,6 +695,13 @@ def from_schema_convention(schema: Schema, weight_features_name_suffix: str = "_
 
         return seq_combiners
 
+    def get_config(self):
+        config = super().get_config()
+        config["axis"] = self.axis
+        config["weight_feature_name"] = self.weight_feature_name
+
+        return config
+
 
 @dataclass
 class EmbeddingOptions: