[jit] Allow custom class functions to be traced in invokeScriptMethod…

…FromPython(). (pytorch#67380)

Summary: Pull Request resolved: pytorch#67380

Test Plan: eyes

Reviewed By: tugsbayasgalan

Differential Revision: D31975656

fbshipit-source-id: 47c8c9854899e9fed5a635f88470711dc4c95970

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)