Deprecate decorating classes with torch.no_grad and similar (#89522)

albanD · pytorchmergebot · commit 347a7d97a5e4 · 2022-11-23T16:51:42.000Z
Fixes #89450 I would have completely removed it but I don't think this is particularly urgent and there are some use of it in the wild: https://github.com/search?q=%2Ftorch%5C.no_grad%5C%28%5C%29%5Cnclass%2F&type=code So we might as well take one release to do it. Pull Request resolved: #89522 Approved by: https://github.com/lezcano, https://github.com/soulitzer, https://github.com/janeyx99
diff --git a/test/test_autograd.py b/test/test_autograd.py
@@ -63,6 +63,30 @@ def graph_desc(fn):
 
 class TestAutograd(TestCase):
 
+    def test_grad_mode_class_decoration(self):
+        # Decorating class is deprecated and should not be used
+        with self.assertWarnsRegex(UserWarning, "Decorating classes is deprecated"):
+            @torch.no_grad()
+            class Foo():
+                pass
+
+        # Decorating functions or methods is fine though
+        with warnings.catch_warnings(record=True) as w:
+            @torch.no_grad()
+            def foo():
+                pass
+
+            class Foo2():
+                @torch.no_grad()
+                def __init__(self):
+                    pass
+
+                @torch.no_grad()
+                def foo(self):
+                    pass
+
+        self.assertEqual(len(w), 0)
+
     def test_tensor_grad_warnings(self):
         dummy = torch.empty(1)
 
diff --git a/torch/autograd/grad_mode.py b/torch/autograd/grad_mode.py
@@ -2,6 +2,7 @@
 import torch
 import functools
 import inspect
+import warnings
 from typing import Any, Callable, TypeVar, cast
 
 __all__ = ['no_grad', 'enable_grad', 'set_grad_enabled',
@@ -18,6 +19,12 @@ class _DecoratorContextManager:
     """Allow a context manager to be used as a decorator"""
 
     def __call__(self, func: F) -> F:
+        if inspect.isclass(func):
+            warnings.warn("Decorating classes is deprecated and will be disabled in "
+                          "future versions. You should only decorate functions or methods. "
+                          "To preserve the current behavior of class decoration, you can "
+                          "directly decorate the `__init__` method and nothing else.")
+
         if inspect.isgeneratorfunction(func):
             return self._wrap_generator(func)
 
