/
glebestraikh
/
adept
Обзор
Документация
Войти
/
glebestraikh
/
adept
Код
Запросы
0
Задачи
Вики
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
python/cpp_bindings/batchnorm2d_py.cpp
26 строк
925 B
kolkir
Add python bindings for device configuration
19 июн 2025, 00:03
19 июн 2025, 00:03
c00b85c
Код
Авторство
О чём код?
#include <pybind11/pybind11.h> #include <pybind11/stl.h> #include <adept/nn/batchnorm2d.hpp> #include <sstream> namespace py = pybind11; using namespace adept; void bind_batchnorm2d(py::module& m) { py::class_<BarchNorm2dImpl, Module, std::shared_ptr<BarchNorm2dImpl>>(m, "BatchNorm2d") .def(py::init<index_t, float32_t, float32_t, device_t, dtype_t, index_t>(), py::arg("features"), py::arg("eps") = 1e-05f, py::arg("momentum") = 0.1f, py::arg("device") = device_t::CPU, py::arg("dtype") = dtype_t::Float32, py::arg("device_id") = 0) .def("parameters", &BarchNorm2dImpl::parameters) .def("forward", &BarchNorm2dImpl::forward) .def("__call__", [](BarchNorm2dImpl& self, Variable& input) { return self.forward(input); }) .def("__repr__", [](const BarchNorm2dImpl& v) { std::stringstream buf; buf << v; return buf.str(); }); }