fix axis in shift_data

google · Dec 18, 2023 · 06d079e · 06d079e
1 parent 01d092a
commit 06d079e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/MaxText/input_pipeline.py b/MaxText/input_pipeline.py
@@ -222,7 +222,7 @@ def preprocessing_pipeline_pygrain(
 
   # Shift inputs for teacher-forced training
   if shift:
-    operations.append(pygrain.MapOperation(map_function=pygrain_operations.ShiftData(axis=0,segmented=pack_examples)))  
+    operations.append(pygrain.MapOperation(map_function=pygrain_operations.ShiftData(axis=1,segmented=pack_examples)))  
 
   index_sampler = pygrain.IndexSampler(
     num_records=len(dataset),