fix pylint bug

AI-Hypercomputer · Oct 3, 2023 · 87e5769 · 87e5769
1 parent dba9845
commit 87e5769
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/MaxText/layers.py b/MaxText/layers.py
@@ -25,7 +25,6 @@
 from jax.sharding import Mesh
 from jax.sharding import PartitionSpec as P
 
-import dataclasses
 import functools
 import operator
 from typing import Any, Callable, Iterable, Optional, Sequence, Tuple, Union
@@ -656,14 +655,14 @@ class Embed(nn.Module):
     dtype: the dtype of the embedding vectors (default: float32).
     embedding_init: embedding initializer.
   """
+  # pylint: disable=attribute-defined-outside-init
   config: Config
   num_embeddings: int
   features: int
   cast_input_dtype: Optional[DType] = None
   dtype: DType = jnp.float32
   attend_dtype: Optional[DType] = None
   embedding_init: Initializer = default_embed_init
-  embedding: Array = dataclasses.field(init=False)
 
   def setup(self):
     self.embedding = self.param(