SeanLee97 / AnglE

Train and Infer Powerful Sentence Embeddings with AnglE | 🔥 SOTA on STS and MTEB Leaderboard
https://arxiv.org/abs/2309.12871
MIT License
398 stars 30 forks source link

Added argument to explicitly pass device to class constructor #20

Closed Asad-Ismail closed 6 months ago

Asad-Ismail commented 6 months ago

It will override set_device method if device is passed if not passed it will use set_device method valid devices just like set_device are "cpu", "mps" and "cuda" example usage model = AnglE.from_pretrained(instructor_model_name, pooling_strategy='cls',device="cpu")

SeanLee97 commented 6 months ago

@Asad-Ismail , many thanks for your contribution! I've merged it into the main branch and published it to pypi: https://pypi.org/project/angle-emb/

Asad-Ismail commented 6 months ago

Awesome thankyou @SeanLee97 !