tweak

2024-10-25 20:04:01 +03:00 · 2024-10-25 20:04:01 +03:00 · e66735527c
commit e66735527c
parent 2e22529c99
2 changed files with 1 additions and 2 deletions
--- a/mochi_preview/dit/joint_model/mod_rmsnorm.py
+++ b/mochi_preview/dit/joint_model/mod_rmsnorm.py
@ -18,6 +18,5 @@ class ModulatedRMSNorm(torch.autograd.Function):

        return x_modulated.type_as(x)

-@torch.compiler.disable()
 def modulated_rmsnorm(x, scale, eps=1e-6):
    return ModulatedRMSNorm.apply(x, scale, eps)