Updates the test cases to use thunderfx (#1553)

Lightning-AI · Dec 13, 2024 · 84ba13c · 84ba13c
1 parent d59f9ab
commit 84ba13c
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 72 deletions.
diff --git a/thunder/tests/framework.py b/thunder/tests/framework.py
@@ -25,7 +25,7 @@
 import thunder.core.utils as utils
 
 from thunder.core.trace import TraceCtx, detached_trace
-from thunder.dynamo import ThunderCompiler
+from thunder.dynamo import thunderfx
 
 import thunder
 
@@ -255,7 +255,7 @@ class DynamoThunderTestExecutor(TestExecutor):
     supported_dtypes = (datatypes.dtype,)
 
     def make_callable(self, fn, **kwargs):
-        return torch.compile(backend=ThunderCompiler(**kwargs))(fn)
+        return thunderfx(fn, **kwargs)
 
 
 # TODO Refactor these executors into the actual executor (sub)modules