Potentially fix module not found

nod-ai · Nov 22, 2024 · 8727221 · 8727221
1 parent 40fbd1a
commit 8727221
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sharktank/tests/layers/mmdit_test.py b/sharktank/tests/layers/mmdit_test.py
@@ -95,7 +95,7 @@ def forward(self, img, txt, vec, pe):
 
         @fxb.export_program(name="mmdit", args=(img, txt, vec, rot), strict=False)
         def _(model, img, txt, vec, rot) -> torch.Tensor:
-            return mod(img, txt, vec, rot)
+            return model(img, txt, vec, rot)
 
         output = aot.export(fxb)
         output.verify()