From c047a3c9b89ca70a5711aa6b9a9fbe3f384bb16e Mon Sep 17 00:00:00 2001 From: Andrei Ivanov Date: Wed, 2 Oct 2024 13:11:40 -0700 Subject: [PATCH 1/3] Fixing warning generated by the use of `torch.cuda.amp.autocast*` functions. --- 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..ff2224318857 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 3ccc5088e6bc9b4de786a5330531a986e8f14a1e Mon Sep 17 00:00:00 2001 From: Andrei Ivanov Date: Wed, 2 Oct 2024 13:46:51 -0700 Subject: [PATCH 2/3] 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 ff2224318857..b7183ec230bc 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() ) From 57e222772e6102c69a24bd24c38a2c4be8e68557 Mon Sep 17 00:00:00 2001 From: Andrei Ivanov Date: Wed, 2 Oct 2024 15:11:34 -0700 Subject: [PATCH 3/3] Fixing lint problems --- python/dgl/backend/pytorch/sparse.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/dgl/backend/pytorch/sparse.py b/python/dgl/backend/pytorch/sparse.py index b7183ec230bc..cea390020a02 100644 --- a/python/dgl/backend/pytorch/sparse.py +++ b/python/dgl/backend/pytorch/sparse.py @@ -154,7 +154,7 @@ def _cast_if_autocast_enabled(*args): if not th.is_autocast_enabled(): return args else: - return th.amp.autocast_mode._cast( + return th.amp.autocast_mode._cast( args, "cuda", th.get_autocast_gpu_dtype() )