Skip to content

Commit f5c6dc8

Browse files
authored
fix: small typing issue (#329)
1 parent 048d1a6 commit f5c6dc8

1 file changed

Lines changed: 2 additions & 2 deletions

File tree

model2vec/train/base.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,8 @@
1515
from torch.nn.utils.rnn import pad_sequence
1616
from tqdm import trange
1717

18-
from model2vec import StaticModel
1918
from model2vec.inference import StaticModelPipeline
19+
from model2vec.model import PathLike, StaticModel
2020
from model2vec.train.dataset import TextDataset
2121
from model2vec.train.utils import (
2222
get_probable_pad_token_id,
@@ -140,7 +140,7 @@ def _initialize(self) -> None:
140140
@classmethod
141141
def from_pretrained(
142142
cls: type[ModelType],
143-
path: str = "minishlab/potion-base-32m",
143+
path: PathLike = "minishlab/potion-base-32m",
144144
*,
145145
token: str | None = None,
146146
**kwargs: Any,

0 commit comments

Comments
 (0)