Upload modeling_fast_esmfold.py with huggingface_hub

modeling_fast_esmfold.py

@@ -1079,9 +1079,21 @@ class FastEsmForProteinFolding(EsmForProteinFolding):
             Dict with "losses" key containing per-step MLM loss values
         """
         self._ensure_ttt_ready()
+        # TTT requires fp32 for stable gradient computation. ESMFold typically
+        # runs the backbone in fp16, but small LoRA updates vanish in half precision.
+        esm_dtype = next(self.esm.parameters()).dtype
+        if esm_dtype != torch.float32:
+            self.esm.float()
+            self.mlm_head.float()
         if self._uses_lora:
-            return self._lora_ttt(seq)
-        return self._legacy_ttt(seq)
+            result = self._lora_ttt(seq)
+        else:
+            result = self._legacy_ttt(seq)
+        # Restore original dtype (backbone back to fp16 for inference)
+        if esm_dtype != torch.float32:
+            self.esm.to(esm_dtype)
+            self.mlm_head.to(esm_dtype)
+        return result
 
     # ---- High-Level API ----