search

weaviate / t2v-transformers-models

This is the repo for the container that holds the models for the text2vec-transformers module
BSD 3-Clause "New" or "Revised" License
39 stars 27 forks source link

Sentence Transformers T5 GPU issue #45

Closed trengrj closed 1 year ago

trengrj commented 1 year ago

When running on a NVIDIA GPU the below line

https://github.com/weaviate/t2v-transformers-models/blob/4a38a7b4f1d66d2bd91b7db45c043a099a63f462/vectorizer.py#L178

Needs to be changed to

labels = torch.tensor(labels).to(device)

This change requires additional changes to the HFModel class to pass through the GPU device.