PyTorch学习总结(二)——基于torch.utils.ffi的自定义C扩展

---

步骤一 准备好你的C代码

首先，你写好你的C函数。

接下来你可以找到一个模块的forward和backward函数的实现，其主要实现输入相加的功能。

在你的.c文件中，你可以使用#include <TH/TH.h>和#include <THC/THC.h>指令来分别包含TH及THC。

ffi工具可以确保编译器在build的过程中找到它们。

/* src/my_lib.c */
#include <TH/TH.h>

int my_lib_add_forward(THFloatTensor *input1, THFloatTensor *input2,
THFloatTensor *output)
{
    if (!THFloatTensor_isSameSizeAs(input1, input2))
        return 0