__init__.py

# test_ops/__init__.py
import torch
from ._C import *

# 注册操作
def add_one(input):
    return torch.ops.test_ops.add_one(input)

def multiply_by_two(input):
    return torch.ops.test_ops.multiply_by_two(input)

# 导出函数
__all__ = ['add_one', 'multiply_by_two']