Skip to content

Commit fb380d6

Browse files
committed
Fix pylint unused-variable warning in quantizations.py
1 parent c7b9871 commit fb380d6

1 file changed

Lines changed: 3 additions & 1 deletion

File tree

src/maxtext/layers/quantizations.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1014,6 +1014,8 @@ def _wrap(self, f, name=None):
10141014
import transformer_engine.jax # pylint: disable=import-outside-toplevel # pytype: disable=import-error
10151015
from transformer_engine.common import recipe # pylint: disable=import-outside-toplevel # pytype: disable=import-error
10161016

1017+
default_recipe = self._recipe
1018+
10171019
class TEWrapper(transformer_engine.jax.flax.module.TransformerEngineBase):
10181020
"""Wrapper module for TransformerEngine quantization."""
10191021

@@ -1022,7 +1024,7 @@ def generate_quantizer_set(
10221024
postfix: str = "",
10231025
variable_collection: str | None = None,
10241026
quantization_checkpoint_name: str | None = None,
1025-
fp8_recipe: recipe.Recipe | None = None,
1027+
fp8_recipe: recipe.Recipe | None = default_recipe,
10261028
n_groups: int | None = None,
10271029
):
10281030
"""Generates a set of quantizers for TransformerEngine."""

0 commit comments

Comments
 (0)