From 4a976adcd79ff1c03dd8263700843def238b9040 Mon Sep 17 00:00:00 2001
From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com>
Date: Sun, 14 Sep 2025 11:29:39 +0800
Subject: [PATCH] Fix fp8 model training

---
 comfy/ops.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/comfy/ops.py b/comfy/ops.py
index 2cc9bbc27..e7a0cf8ab 100644
--- a/comfy/ops.py
+++ b/comfy/ops.py
@@ -336,12 +336,13 @@ class fp8_ops(manual_cast):
             return None
 
         def forward_comfy_cast_weights(self, input):
-            try:
-                out = fp8_linear(self, input)
-                if out is not None:
-                    return out
-            except Exception as e:
-                logging.info("Exception during fp8 op: {}".format(e))
+            if not self.training:
+                try:
+                    out = fp8_linear(self, input)
+                    if out is not None:
+                        return out
+                except Exception as e:
+                    logging.info("Exception during fp8 op: {}".format(e))
 
             weight, bias = cast_bias_weight(self, input)
             return torch.nn.functional.linear(input, weight, bias)