mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
Previously, the doc's suggested `from torch._library.triton import wrap_triton, triton_op` doesn't work because wrap_triton is not imported in torch/_library/__init__.py but `from torch.library import wrap_triton` works. This PR imports wrap_triton and fix the doc. Pull Request resolved: https://github.com/pytorch/pytorch/pull/144882 Approved by: https://github.com/zou3519
7 lines
269 B
Python
7 lines
269 B
Python
import torch._library.autograd
|
|
import torch._library.fake_impl
|
|
import torch._library.simple_registry
|
|
import torch._library.utils
|
|
from torch._library.fake_class_registry import register_fake_class
|
|
from torch._library.triton import capture_triton, triton_op, wrap_triton
|