lexir / check_device.py
irinaqqq's picture
Upload folder using huggingface_hub
6a02b16 verified
raw
history blame contribute delete
855 Bytes
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()