Addind a comment to get_input_bld

Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

Author: mingxu1067

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