File size: 855 Bytes
6a02b16 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 | import torch
from sentence_transformers import SentenceTransformer
def main():
cuda_available = torch.cuda.is_available()
print("torch.cuda.is_available:", cuda_available)
if cuda_available:
print("cuda device count:", torch.cuda.device_count())
for i in range(torch.cuda.device_count()):
print(f"cuda:{i} name:", torch.cuda.get_device_name(i))
print("torch version:", torch.__version__)
print("cuda version:", torch.version.cuda)
device = "cuda"
else:
device = "cpu"
model = SentenceTransformer("paraphrase-multilingual-mpnet-base-v2", device=device)
print("sentence-transformers device:", model.device)
x = model.encode(["тест"], convert_to_numpy=True)
print("encode ok, shape:", x.shape)
if __name__ == "__main__":
main()
|