Update app.py

app.py

@@ -181,7 +181,7 @@ def process_hdf5_file(uploaded_file, percentage_idx):
         inference = load_module_from_path("inference", inference_path)
 
         # Step 4: Load the model from lwm_model module
-        device = 'cpu'
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
         print(f"Loading the LWM model on {device}...")
         model = lwm_model.LWM.from_pretrained(device=device)