fix typo

mosaicml · Nov 16, 2023 · 6791f3d · 6791f3d
1 parent 2f8afa2
commit 6791f3d
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tests/test_mpt_gen.py b/tests/test_mpt_gen.py
@@ -147,7 +147,8 @@ def test_mpt_generate_callback(attn_impl: str, use_alibi: bool,
 @pytest.mark.parametrize('attn_impl', ['triton', 'torch'])
 @pytest.mark.parametrize('use_alibi', [True, False])
 def test_mpt_generate_callback_not_tied(
-        attn_impl: str, build_tiny_mpt: Callable[..., ComposerMPTCausalLM],
+        use_alibi: bool, attn_impl: str,
+        build_tiny_mpt: Callable[..., ComposerMPTCausalLM],
         tiny_ft_dataloader: DataLoader):
     device = get_device('gpu')
 
@@ -157,7 +158,7 @@ def test_mpt_generate_callback_not_tied(
         attn_config={
             'attn_impl': attn_impl,
             'attn_uses_sequence_id': False,
-            'alibi': True,
+            'alibi': use_alibi,
         },
     )
     model = device.module_to_device(model)