replace boolean causal mask for int64 causal mask

neuralmagic · Jul 26, 2023 · 91118e6 · 91118e6
1 parent e324cdc
commit 91118e6
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/src/deepsparse/transformers/utils/helpers.py b/src/deepsparse/transformers/utils/helpers.py
@@ -58,7 +58,7 @@ def pad_to_fixed_length(
 def create_causal_mask(
     input_ids: Union[numpy.ndarray, List[int]],
     attention_mask: Union[numpy.ndarray, List[int]],
-    dtype: numpy.dtype = numpy.bool_,
+    dtype: numpy.dtype = numpy.int64,
 ) -> numpy.ndarray:
     """
     Compute a causal from a set of module inputs.

diff --git a/tests/deepsparse/transformers/utils/helpers.py b/tests/deepsparse/transformers/utils/helpers.py
@@ -84,6 +84,4 @@
 )
 def test_create_causal_mask(input_ids, attention_mask, expected_causal_mask):
     causal_mask = create_causal_mask(input_ids, attention_mask)
-    assert numpy.array_equal(
-        causal_mask, expected_causal_mask[None, None, ...].astype(numpy.bool_)
-    )
+    assert numpy.array_equal(causal_mask, expected_causal_mask[None, None, ...])