File size: 179 Bytes
789eef1 |
1 2 3 4 5 6 7 8 |
#include <torch/extension.h>
#include "ssim.h"
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("fusedssim", &fusedssim);
m.def("fusedssim_backward", &fusedssim_backward);
}
|