checkVRAM

2da2e7a7 · Мазур Грета Евгеньевна · fee925be · 2da2e7a7
Commit 2da2e7a7 authored 3 weeks ago by Мазур Грета Евгеньевна
Hide whitespace changes
Inline Side-by-side

Showing

with 2 additions and 0 deletions
+2 -0
--- a/grustno.py
+++ b/grustno.py
@@ -9,6 +9,8 @@ import torch
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(device)
 # model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=4).to(device)
+if torch.cuda.is_available():
+    print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9} GB")
 import pandas as pd