Addind a comment to get_input_bld

Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
google · Mar 5, 2024 · e3e785c · e3e785c
1 parent 0bd4a53
commit e3e785c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/praxis/contrib/gpu/scripts_gpu/te_helper.py b/praxis/contrib/gpu/scripts_gpu/te_helper.py
@@ -50,6 +50,7 @@ def set_layer_params_to_stack_transformer(stacked_transformer_obj, layer_p, laye
 
     @staticmethod
     def get_input_bld(original_bld, batch_axes, mdl_axis):
+        # This is used to specify the sharding pattern of inputs to TransformerLayers.
         raise NotImplementedError
 
     @staticmethod
@@ -89,7 +90,7 @@ def set_layer_params_to_stack_transformer(stacked_transformer_obj, layer_p, laye
         return layer_p
 
     @staticmethod
-    def get_input_bld(original_bld, _, _):
+    def get_input_bld(original_bld, *_):
         return original_bld
 
     @staticmethod