mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
9e8d2b6de2
The user provides a `setup_context` and a `backward_function`. These get put into a torch.autograd.Function that gets registered as the custom op's autograd implementation. Test Plan: - we update custom ops in the custom_op_db to use the new register_autograd API. Pull Request resolved: https://github.com/pytorch/pytorch/pull/123110 Approved by: https://github.com/albanD ghstack dependencies: #123108, #123109
6 lines
201 B
Python
6 lines
201 B
Python
import torch._library.abstract_impl
|
|
import torch._library.autograd
|
|
import torch._library.simple_registry
|
|
import torch._library.utils
|
|
|
|
from torch._library.fake_class_registry import register_fake_class
|