import torch

# Check if CUDA (GPU) is available, otherwise use CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load the checkpoint
lora_path = "./lora_checkpoint.pt"  # Path to the checkpoint file
checkpoint = torch.load(lora_path, map_location=device)

# Print the checkpoint keys to inspect the structure
print("Checkpoint keys:", checkpoint.keys())