pytorch/torch/csrc/jit/script/module_python.h
Zachary DeVito 5b87049c66 remove uses of std::shared_ptr<Module> (#21934)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/21934
ghimport-source-id: e64ab9096f43749ead3ac5567675b815da295664

Test Plan: Imported from OSS

Differential Revision: D15892401

Pulled By: zdevito

fbshipit-source-id: 6424139206593ff944556c69d8a54723884eacaf
2019-06-25 13:24:38 -07:00

23 lines
488 B
C++

#pragma once
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <torch/csrc/jit/script/module.h>
namespace py = pybind11;
namespace torch {
namespace jit {
namespace script {
inline c10::optional<Module> as_module(const py::object& obj) {
if (py::isinstance(
obj, py::module::import("torch.jit").attr("ScriptModule"))) {
return py::cast<Module>(obj.attr("_c"));
}
return c10::nullopt;
}
} // namespace script
} // namespace jit
} // namespace torch