Call grad_mode.py context managers as decorators

[jvmncs]

Extends the context manager classes torch.no_grad, torch.enable_grad, and torch.set_grad_mode to function as decorators, so that users can wrap functions that will never require a call to .backward() downstream. I've modified the docs to reflect this change, and I've also added tests for the new functionality in each mode's respective test in test_autograd.py.

I also didn't find a unit test specifically for torch.enable_grad. Assuming that's intended, unless my ctrl+f missed it.

[ssnl]

LGTM if test passes

[ssnl]

@pytorchbot retest this please

[ssnl]

@pytorchbot retest this please

[ssnl]

@pytorchbot retest this please

[ssnl]

@pytorchbot test this please

[ssnl]

test failure looks legit:

19:56:37 FAIL: test_set_grad_enabled (test_autograd.TestAutograd)
19:56:37 ----------------------------------------------------------------------
19:56:37 Traceback (most recent call last):
19:56:37   File "test_autograd.py", line 1812, in test_set_grad_enabled
19:56:37     self.assertTrue(y.requires_grad)
19:56:37 AssertionError: False is not true

[jvmncs]

ah, indeed. going to have to change up the set_grad_enabled class. nice catch

[jvmncs]

so, question. in order to maintain the current behavior, e.g. to use torch.set_grad_enabled(False) imperatively, I need to be able to save the input to set_grad_enabled as an attribute, then change enter to set the grad mode to that attribute's value:

def __init__(self, mode):
    self.prev = torch.is_grad_enabled()
    torch._C.set_grad_enabled(mode)
    self.mode = mode

def __enter__(self):
    torch.set_grad_enabled(self.mode)

def __exit__(self, *args):
    torch.set_grad_enabled(self.prev)
    return False

however, if I do that, then any time it's instantiated, including when it wraps a function, the underlying grad mode will be changed. this seems like unwanted behavior, and maybe it's best to not use this one as a decorator? open to other suggestions.

[apaszke]

I think it's ok to forbid using set_grad_mode as a decorator, since its generally useful when run-time values are changing, but function definition usually happens only once at the beginning of the program.