以下代码将手动把tanh算子fallback到torch_npu执行,与torch_npu的配套使用时请参见配套Torch_NPU使用安装torch_npu环境。
示例代码如下所示:
import torch
import torch.nn as nn
import torch_npu
import mindietorch
mindietorch.set_device(device_id)
class Test(nn.Module):
def forward(self, x):
x = torch.ops.aten.relu.default(x)
x= torch.ops.aten.tanh.default(x)
out = torch.ops.aten.sigmoid.default(x)
return out
shape = (2, 2)
input = torch.randn(shape)
model = Test().to("npu")
backend_kwargs = {
"torch_executed_ops": [torch.ops.aten.tanh.default],
"min_block_size": 1,
}
opt_model = torch.compile(model, backend="mindie", options=backend_kwargs)
device_id = 0
npu_input = input.to("npu")
infer_ret = compiled_model(npu_input)[0].to("cpu")