input data type in fuji

yangdsh · Mar 12, 2024 · 47b21f0 · 47b21f0
1 parent d4265dc
commit 47b21f0
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/axlearn/experiments/text/gpt/fuji.py b/axlearn/experiments/text/gpt/fuji.py
@@ -16,6 +16,7 @@
     RepeatedTransformerLayer,
     RoFormerQKVLinear,
 )
+from axlearn.common.utils import DataPartitionType
 from axlearn.common.embedding import TransformerTextEmbeddings
 from axlearn.common.layers import RMSNorm
 from axlearn.experiments.text.gpt.common import STEP_DTYPE, learner_config, mesh_shape_from_axes