Update app.py

app.py

@@ -8,6 +8,8 @@ from transformers import AutoTokenizer
 from safetensors.torch import load_file as load_sft
 from huggingface_hub import snapshot_download
 
+torch.set_default_dtype(torch.float32)  
+
 # ===============================================
 # Default config (from your training notes)
 # ===============================================
@@ -60,6 +62,9 @@ class AttnBlock(nn.Module):
         return Qh2, Kh2
 
     def forward(self, x, rope, radius):
+        if x.dtype != self.norm1.weight.dtype:
+        x = x.to(self.norm1.weight.dtype)
+        
         h = self.norm1(x)
         B, S, E = h.shape
         cos, sin = rope
@@ -382,11 +387,15 @@ def load_model(source: str):
     else:
         model.load_state_dict(state, strict=True)
     
-    # 🔧 ensure everything is float32 (fixes mixed-dtype LayerNorm on CPU)
+    # ✅ hard-cast ALL params & buffers to float32 (handles weights-only .pt that saved as float64)
     model = model.to(torch.float32)
-    for bname, buf in model.named_buffers():
-        if buf.dtype.is_floating_point:
-            setattr(model, bname, buf.float())
+    with torch.no_grad():
+        for p in model.parameters():
+            if p.dtype.is_floating_point:
+                p.data = p.data.float()
+        for _, buf in model.named_buffers():
+            if buf.dtype.is_floating_point:
+                buf.data = buf.data.float()
     
     model.eval()
     return model, tokenizer, conf["radius"]