import torch

# Check if GPU is available
if torch.cuda.is_available():
    device = torch.device("cuda")
    print("GPU is available. Using GPU.")
else:
    device = torch.device("cpu")
    print("GPU not available. Using CPU.")

# Example model loading
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "cerebras/btlm-3b-8k-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Move model to the appropriate device
model.to(device)