Expand _skip_layerwise_casting_patterns list

2026-03-05 20:16:19 +08:00
parent 1f6b1e2c08
commit da901e7368
1 changed files with 9 additions and 2 deletions
@@ -536,7 +536,14 @@ class HeliosTransformer3DModel(
    """

    _supports_gradient_checkpointing = True
-    _skip_layerwise_casting_patterns = ["patch_embedding", "condition_embedder", "norm"]
+    _skip_layerwise_casting_patterns = [
+        "patch_embedding",
+        "patch_short",
+        "patch_mid",
+        "patch_long",
+        "condition_embedder",
+        "norm",
+    ]
    _no_split_modules = ["HeliosTransformerBlock", "HeliosOutputNorm"]
    _keep_in_fp32_modules = [
        "time_embedder",
@@ -808,4 +815,4 @@ class HeliosTransformer3DModel(
        if not return_dict:
            return (output,)

-        return Transformer2DModelOutput(sample=output)
+        return Transformer2DModelOutput(sample=output)