From 323f18c8c4e63040d9cc25a9562dc96f97147c8a Mon Sep 17 00:00:00 2001 From: Andrei Ivanov Date: Thu, 30 Jan 2025 10:48:24 -0800 Subject: [PATCH 1/2] Resolving warnings related to `amp.autocast` --- python/dgl/backend/pytorch/sparse.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/dgl/backend/pytorch/sparse.py b/python/dgl/backend/pytorch/sparse.py index 67034ad99579..a7d0888f3c22 100644 --- a/python/dgl/backend/pytorch/sparse.py +++ b/python/dgl/backend/pytorch/sparse.py @@ -145,7 +145,7 @@ def __exit__(self, *args, **kargs): # and do it only in a nested autocast context. def _disable_autocast_if_enabled(): if th.is_autocast_enabled(): - return th.cuda.amp.autocast(enabled=False) + return th.amp.autocast('cuda', enabled=False) else: return empty_context() @@ -154,8 +154,8 @@ def _cast_if_autocast_enabled(*args): if not th.is_autocast_enabled(): return args else: - return th.cuda.amp.autocast_mode._cast( - args, th.get_autocast_gpu_dtype() + return th.amp.autocast_mode._cast( + args, 'cuda', th.get_autocast_gpu_dtype() ) From 7ccf30bab44fda135f52aed365247dc373a6b25b Mon Sep 17 00:00:00 2001 From: Andrei Ivanov Date: Thu, 30 Jan 2025 11:06:57 -0800 Subject: [PATCH 2/2] Fixing lint problems --- python/dgl/backend/pytorch/sparse.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/dgl/backend/pytorch/sparse.py b/python/dgl/backend/pytorch/sparse.py index a7d0888f3c22..cea390020a02 100644 --- a/python/dgl/backend/pytorch/sparse.py +++ b/python/dgl/backend/pytorch/sparse.py @@ -145,7 +145,7 @@ def __exit__(self, *args, **kargs): # and do it only in a nested autocast context. def _disable_autocast_if_enabled(): if th.is_autocast_enabled(): - return th.amp.autocast('cuda', enabled=False) + return th.amp.autocast("cuda", enabled=False) else: return empty_context() @@ -155,7 +155,7 @@ def _cast_if_autocast_enabled(*args): return args else: return th.amp.autocast_mode._cast( - args, 'cuda', th.get_autocast_gpu_dtype() + args, "cuda", th.get_autocast_gpu_dtype() )