From f20614af21b6113382cc17e85daf22febd1daff4 Mon Sep 17 00:00:00 2001 From: Zhengxu Chen Date: Wed, 27 Oct 2021 16:37:24 -0700 Subject: [PATCH] [jit] Allow custom class functions to be traced in invokeScriptMethodFromPython(). (#67380) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/67380 Test Plan: eyes Reviewed By: tugsbayasgalan Differential Revision: D31975656 fbshipit-source-id: 47c8c9854899e9fed5a635f88470711dc4c95970 --- torch/csrc/jit/python/pybind_utils.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch/csrc/jit/python/pybind_utils.h b/torch/csrc/jit/python/pybind_utils.h index c2133d69cb3f0..f277b031e4cf7 100644 --- a/torch/csrc/jit/python/pybind_utils.h +++ b/torch/csrc/jit/python/pybind_utils.h @@ -979,7 +979,7 @@ inline py::object runAndInsertCall( // and then run the callee with tracing disabled. // Get the graph `Value`s that represent the input IValues - auto inputs = last(stack, toGraphFunction(callee).graph()->inputs().size()); + auto inputs = last(stack, callee.num_inputs()); auto input_values = fmap(inputs, [](const IValue& v) { return tracer::getValueTrace(v); }); TORCH_INTERNAL_ASSERT(callee.getSchema().returns().size() == 1)