make has_bundled_input work for flatbuffer (#76854)

Summary: title

Test Plan: unit test

Differential Revision: D36120947

Pull Request resolved: https://github.com/pytorch/pytorch/pull/76854
Approved by: https://github.com/Jack-Khuu
This commit is contained in:
Han Qi
2022-05-09 23:04:08 +00:00
committed by PyTorch MergeBot
parent b4fa1e88be
commit 41ff6f8c49

View File

@ -12,6 +12,7 @@
#include <torch/csrc/jit/mobile/code.h>
#include <torch/csrc/jit/mobile/compatibility/backport.h>
#include <torch/csrc/jit/mobile/compatibility/model_compatibility.h>
#include <torch/csrc/jit/mobile/file_format.h>
#include <torch/csrc/jit/mobile/import.h>
#include <torch/csrc/jit/mobile/module.h>
#include <torch/csrc/jit/operator_upgraders/upgraders.h>
@ -2244,6 +2245,17 @@ void initJitScriptBindings(PyObject* module) {
return py::isinstance<Object>(obj);
});
m.def("_get_file_format", [](const std::string& path) {
switch (getFileFormat(path)) {
case FileFormat::FlatbufferFileFormat:
return "flatbuffer";
case FileFormat::ZipFileFormat:
return "zipfile";
default:
return "invalid";
}
});
initScriptDictBindings(module);
initScriptListBindings(module);
}