Skip to content

Commit

Permalink
Fix docstring
Browse files Browse the repository at this point in the history
Activation functions are passed as callable to Megatron-LM core, so
mentioning it here is incorrect.
  • Loading branch information
janEbert committed May 17, 2024
1 parent 44f2609 commit 3943cf8
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion nemo/collections/nlp/modules/common/megatron/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -160,7 +160,7 @@ def attention_mask_func(attention_scores, attention_mask):

def get_te_activation(activation: str) -> str:
"""Return the string for the desired activation converted for usage
with TransformerEngine and/or Megatron-LM core.
with TransformerEngine.
Log an error if it cannot be converted.
"""
Expand Down

0 comments on commit 3943cf8

Please sign in to comment.