Fix type hint of StaticEmbedding.__init__ (#3196)

UKPLab · Jan 25, 2025 · db6ce94 · db6ce94
1 parent dfef1d6
commit db6ce94
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/sentence_transformers/models/StaticEmbedding.py b/sentence_transformers/models/StaticEmbedding.py
@@ -19,7 +19,7 @@ class StaticEmbedding(nn.Module):
     def __init__(
         self,
         tokenizer: Tokenizer | PreTrainedTokenizerFast,
-        embedding_weights: np.array | torch.Tensor | None = None,
+        embedding_weights: np.ndarray | torch.Tensor | None = None,
         embedding_dim: int | None = None,
         **kwargs,
     ) -> None:
@@ -30,7 +30,7 @@ def __init__(
         Args:
             tokenizer (Tokenizer | PreTrainedTokenizerFast): The tokenizer to be used. Must be a fast tokenizer
                 from ``transformers`` or ``tokenizers``.
-            embedding_weights (np.array | torch.Tensor | None, optional): Pre-trained embedding weights.
+            embedding_weights (np.ndarray | torch.Tensor | None, optional): Pre-trained embedding weights.
                 Defaults to None.
             embedding_dim (int | None, optional): Dimension of the embeddings. Required if embedding_weights
                 is not provided. Defaults to None.