mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-15 21:00:47 +00:00
f2f6e6c0e8
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/10759 Adding a basic registry pattern to pybindstate so that we can have separate 'cc' files register module updates. This is substantially cleaner than using multiple pybind modules (which have been known to cause bugs) Reviewed By: bddppq Differential Revision: D9441878 fbshipit-source-id: af9e9e98385e92b58ca50e935678328c62684d8e
33 lines
958 B
C++
33 lines
958 B
C++
#pragma once
|
|
|
|
#include <pybind11/pybind11.h>
|
|
#include "caffe2/core/registry.h"
|
|
|
|
namespace caffe2 {
|
|
namespace python {
|
|
|
|
namespace py = pybind11;
|
|
|
|
struct PybindAddition {
|
|
PybindAddition() {}
|
|
PybindAddition(py::module&) {}
|
|
virtual ~PybindAddition(){};
|
|
};
|
|
|
|
CAFFE_DECLARE_REGISTRY(PybindAdditionRegistry, PybindAddition, py::module&);
|
|
|
|
#define REGISTER_PYBIND_ADDITION(funcname) \
|
|
namespace { \
|
|
struct funcname##Impl : public PybindAddition { \
|
|
funcname##Impl(py::module& m) { \
|
|
funcname(m); \
|
|
} \
|
|
}; \
|
|
CAFFE_REGISTER_CLASS( \
|
|
PybindAdditionRegistry, \
|
|
funcname##Impl, \
|
|
funcname##Impl); \
|
|
}
|
|
|
|
} // namespace python
|
|
} // namespace caffe2
|