[bugfix] Set stacklevel=2 for dgl_warning (dmlc#3816)

Co-authored-by: Jinjing Zhou <[email protected]> Co-authored-by: Minjie Wang <[email protected]>
skorani · Mar 25, 2022 · f73c7a6 · f73c7a6
1 parent 7d41608
commit f73c7a6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/dgl/base.py b/python/dgl/base.py
@@ -40,9 +40,9 @@ def dgl_warning_format(message, category, filename, lineno, line=None):
     else:
         return _default_formatwarning(message, category, filename, lineno, line=None)
 
-def dgl_warning(message, category=DGLWarning, stacklevel=1):
+def dgl_warning(message, category=DGLWarning, stacklevel=2):
     """DGL warning wrapper that defaults to ``DGLWarning`` instead of ``UserWarning`` category."""
-    return warnings.warn(message, category=category, stacklevel=1)
+    return warnings.warn(message, category=category, stacklevel=stacklevel)
 
 warnings.formatwarning = dgl_warning_format