[diffusion] fix: fix torch.compile graph break caused by torch._dynamo.disable (#18336)

2026-02-06 14:48:09 +08:00
parent d0c39bc219
commit f798ab9775
3 changed files with 38 additions and 23 deletions
@@ -319,6 +319,11 @@ class _ScaleResidualNormScaleShift(CustomOp):
            fused_scale_residual_norm_scale_shift,
        )

+        if isinstance(gate, int) and gate != 1:
+            raise ValueError(
+                f"Only gate value of 1 is supported for int type, but got {gate}"
+            )
+
        return fused_scale_residual_norm_scale_shift(
            residual.contiguous(),
            x.contiguous(),