Update app.py

app.py

@@ -589,10 +589,12 @@ def on_prepare_loras_click(
         print("[LoRA] No LoRAs selected, resetting to base model weights")
         try:
             transformer = ledger.transformer()
-            base_weights = {k: v.cpu() for k, v in transformer.state_dict().items()}
+            target_device = next(transformer.parameters()).device
+            
+            # Get base weights and keep them on the same device as transformer
+            base_weights = {k: v.to(target_device) for k, v in transformer.state_dict().items()}
             transformer.load_state_dict(base_weights, strict=False)
-            if torch.cuda.is_available():
-                transformer = transformer.to("cuda")
+            
             current_lora_key = key
             progress(1.0, desc="Done")
             return "✓ Reset to base model (no LoRAs active)"
@@ -624,23 +626,20 @@ def on_prepare_loras_click(
     
     try:
         transformer = ledger.transformer()
-        target_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         
-        # Move transformer to CPU for loading (avoids device mismatch)
-        transformer = transformer.to("cpu")
-        torch.cuda.empty_cache()  # Free VRAM from the CPU copy
+        # Determine target device - the transformer should already be on GPU
+        target_device = next(transformer.parameters()).device
         
-        # Load the fused state dict
-        missing, unexpected = transformer.load_state_dict(fused_state, strict=False)
+        # Load fused state dict directly into GPU transformer
+        # Convert CPU tensors to GPU tensors inline, then load
+        fused_state_gpu = {k: v.to(target_device) for k, v in fused_state.items()}
+        
+        missing, unexpected = transformer.load_state_dict(fused_state_gpu, strict=False)
         if missing:
             print(f"[LoRA] Warning: {len(missing)} keys not found in fused state")
         if unexpected:
             print(f"[LoRA] Warning: {len(unexpected)} unexpected keys in fused state")
         
-        # Move transformer to target device (GPU for generation)
-        if target_device.type != "cpu":
-            transformer = transformer.to(target_device)
-        
         current_lora_key = key
         progress(1.0, desc="Done")
         return f"✓ Applied {len(active_loras)} LoRA(s) successfully"