File size: 1,112 Bytes
dc7e099 | 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 27 28 29 30 31 32 33 34 35 36 37 | import torch
import sys
import os
MODEL_PATH = "/workspace/ankahi/punjabi_gemma/ankahi"
print(f"Python version: {sys.version}")
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
print(f"CUDA version: {torch.version.cuda}")
print(f"Device count: {torch.cuda.device_count()}")
for i in range(torch.cuda.device_count()):
print(f"Device {i}: {torch.cuda.get_device_name(i)}")
try:
from transformers import AutoModelForCausalLM, AutoProcessor
print("Transformers imported.")
print("Loading processor...")
processor = AutoProcessor.from_pretrained(MODEL_PATH)
print("Loading model on CPU first to check if weights are readable...")
# This might take a lot of RAM
model = AutoModelForCausalLM.from_pretrained(
MODEL_PATH,
torch_dtype=torch.bfloat16,
low_cpu_mem_usage=True,
device_map="auto"
)
print("Model loaded successfully!")
except Exception as e:
print(f"Caught exception: {e}")
import traceback
traceback.print_exc()
|