Support Torch Compile (#1368)

* support simple compile

* add support for compile

* minor fix

* minor fix

* minor fix
This commit is contained in:
Hong Zhang
2026-03-24 11:19:43 +08:00
committed by GitHub
parent e2a3a987da
commit ae8cb139e8
18 changed files with 59 additions and 2 deletions

View File

@@ -549,6 +549,9 @@ class QwenImageTransformerBlock(nn.Module):
class QwenImageDiT(torch.nn.Module):
_repeated_blocks = ["QwenImageTransformerBlock"]
def __init__(
self,
num_layers: int = 60,