File size: 402 Bytes
d1ceb73 |
1 2 3 4 5 6 7 8 9 10 11 12 13 |
# mypy: allow-untyped-defs
import torch
from torch._ops import OpOverload, OpOverloadPacket
def _register_decomposition(op: OpOverload, graph: torch._C.Graph):
assert not isinstance(
op, OpOverloadPacket
), f"Must pass specific op overload, not overload packet, found {op}"
assert isinstance(op, OpOverload)
torch._C._jit_register_decomposition_for_schema(op._schema, graph)
|