Update app.py

app.py

@@ -1,19 +1,11 @@
-# Inspecting the keys: This will show you the structure of the checkpoint and help you identify the correct components to load into your model.
+import torch
 
-import torch  # Import torch for loading the checkpoint
+# Check if CUDA (GPU) is available, otherwise use CPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Set the path to the LoRA checkpoint file
-lora_path = "./lora_checkpoint.pt"  # Assuming the checkpoint is in the current folder
-
-# Load the checkpoint from the specified path
+# Load the checkpoint
+lora_path = "./lora_checkpoint.pt"  # Path to the checkpoint file
 checkpoint = torch.load(lora_path, map_location=device)
 
-# Print the keys to see the structure
+# Print the checkpoint keys to inspect the structure
 print("Checkpoint keys:", checkpoint.keys())
-
-# If you want to inspect the values of a specific key
-for key in checkpoint.keys():
-    print(f"Key: {key}, Value shape: {checkpoint[key].shape if hasattr(checkpoint[key], 'shape') else 'N/A'}")
-    
-# You can also print the entire checkpoint if it is not too large
-# print(checkpoint)