linting - add raises NotImplError

pytorch · Feb 26, 2024 · ae59079 · ae59079
1 parent 508c73c
commit ae59079
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/torchtrain/models/llama/model.py b/torchtrain/models/llama/model.py
@@ -401,6 +401,9 @@ class TransformerBlock(nn.Module):
         attention_norm (RMSNorm): Layer normalization for attention output.
         ffn_norm (RMSNorm): Layer normalization for feedforward output.
 
+    Raises:
+        NotImplementedError: If norm_type is not rmsnorm or fastlayernorm.
+
     """
 
     def __init__(self, layer_id: int, model_args: ModelArgs):