Files
kernel-lab/tasks/06_pytorch_custom_op/extension_skeleton.py
2026-04-10 13:15:06 +00:00

27 lines
590 B
Python

from __future__ import annotations
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[2]
if str(ROOT) not in sys.path:
sys.path.insert(0, str(ROOT))
import torch
from tools.lab_extension import build_extension
def main() -> None:
ext = build_extension(verbose=True)
if ext is None:
return
print("Extension loaded.")
print("Available torch.ops namespace:", hasattr(torch.ops, "kernel_lab"))
if hasattr(torch.ops, "kernel_lab"):
print("Registered ops:", dir(torch.ops.kernel_lab))
if __name__ == "__main__":
main()