Refactor code for speed and clarity
This commit is contained in:
Glenn Jocher authored and GitHub committed 2024-06-30 12:34:04 +02:00
1 parent 05e59bdec1
commit b9019671f4
13 files changed
+29 -5

No files matched your search

+1
View File
@@ -37,6 +37,7 @@ def smart_inference_mode(torch_1_9=check_version(torch.__version__, "1.9.0")):
"""Applies torch.inference_mode() if torch>=1.9.0, else torch.no_grad() as a decorator for functions."""
def decorate(fn):
"""Applies torch.inference_mode() if torch>=1.9.0, else torch.no_grad() to the decorated function."""
return (torch.inference_mode if torch_1_9 else torch.no_grad)()(fn)
return decorate