mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
b6eb7d7ba0
This allows _broadcast and _broadcast_coalesced to be called on Variables. It also broadens the types accepted by some JIT methods.
40 lines
950 B
C++
40 lines
950 B
C++
#pragma once
|
|
|
|
#include <Python.h>
|
|
#include <ATen/ATen.h>
|
|
#include <pybind11/pybind11.h>
|
|
#include <pybind11/stl.h>
|
|
|
|
#include "torch/csrc/DynamicTypes.h"
|
|
#include "torch/csrc/autograd/python_variable.h"
|
|
|
|
namespace py = pybind11;
|
|
|
|
namespace pybind11 { namespace detail {
|
|
|
|
// handle Tensor <-> at::Tensor conversions
|
|
// Python Variables are unpacked into Tensors
|
|
template <> struct type_caster<at::Tensor> {
|
|
public:
|
|
PYBIND11_TYPE_CASTER(at::Tensor, _("at::Tensor"));
|
|
|
|
bool load(handle src, bool) {
|
|
PyObject* obj = src.ptr();
|
|
if (THPVariable_Check(obj)) {
|
|
value = ((THPVariable*)obj)->cdata.data();
|
|
return true;
|
|
}
|
|
try {
|
|
value = torch::createTensor(obj);
|
|
} catch (std::exception& e) {
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
static handle cast(at::Tensor src, return_value_policy /* policy */, handle /* parent */) {
|
|
return handle(torch::createPyObject(src));
|
|
}
|
|
};
|
|
|
|
}} // namespace pybind11::detail
|