Fix arg name in numerics test

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
NVIDIA · Jan 19, 2024 · 2ceeee7 · 2ceeee7
1 parent 051db0d
commit 2ceeee7
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
@@ -1225,7 +1225,7 @@ def test_transformer_layer_hidden_states_format(dtype, bs, model):
             kv_channels=config.embed,
             apply_residual_connection_post_layernorm=False,
             output_layernorm=False,
-            hidden_states_format="sbhd"
+            attn_input_format="sbhd"
         )
         .to(dtype=dtype)
         .cuda()
@@ -1248,7 +1248,7 @@ def test_transformer_layer_hidden_states_format(dtype, bs, model):
             kv_channels=config.embed,
             apply_residual_connection_post_layernorm=False,
             output_layernorm=False,
-            hidden_states_format="bshd"
+            attn_input_format="bshd"
         )
         .to(dtype=dtype)
         .cuda()