Try different monkey-patch

app.py

@@ -47,16 +47,23 @@ def load_model():
     )
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = torch.load(ckpt_path, map_location=device)
 
-    # 🔧 Monkey-patch for missing PEFT attribute
-    def add_cast_input_dtype_enabled(module):
-        for child in module.children():
-            add_cast_input_dtype_enabled(child)
-        if isinstance(module, torch.nn.Linear) and not hasattr(module, "cast_input_dtype_enabled"):
-            module.cast_input_dtype_enabled = False
+    # Step 1: Create model from scratch
+    model = CustomTransformerModel(CustomTransformerConfig())
 
-    add_cast_input_dtype_enabled(model)
+    # Step 2: Load state_dict from full checkpoint
+    full_model = torch.load(ckpt_path, map_location=device)
+
+    # This handles both full model or just state_dict
+    try:
+        state_dict = full_model.state_dict()
+    except AttributeError:
+        state_dict = full_model  # already a state_dict
+
+    # Step 3: Load weights (might print mismatches)
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    print("Missing keys:", missing)
+    print("Unexpected keys:", unexpected)
 
     model = disable_dropout(model)
     model.to(device)
@@ -64,6 +71,7 @@ def load_model():
     return model
 
 
+
 rng = np.random.default_rng()
 
 # --- Utility Functions ---