diff --git a/docs/conf.py b/docs/conf.py
index 4d62c5650..8a6fabce7 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -45,7 +45,7 @@ def setup(app):
 
         def wrapped(obj, **kwargs):
             import triton
-            if isinstance(obj, triton.code_gen.JITFunction):
+            if isinstance(obj, triton.runtime.JITFunction):
                 obj = obj.fn
             return old(obj)
 
@@ -56,7 +56,7 @@ def setup(app):
 
     def documenter(app, obj, parent):
         import triton
-        if isinstance(obj, triton.code_gen.JITFunction):
+        if isinstance(obj, triton.runtime.JITFunction):
             obj = obj.fn
         return old_documenter(app, obj, parent)