Skip to content

Instantly share code, notes, and snippets.

@dhruvbird
Created November 16, 2021 23:37
Show Gist options
  • Save dhruvbird/fee7cf0c6d2ec2b8f494245e32364c11 to your computer and use it in GitHub Desktop.
Save dhruvbird/fee7cf0c6d2ec2b8f494245e32364c11 to your computer and use it in GitHub Desktop.
Fetch metadata only using PyTorch Lite Model C++ API
#include <torch/csrc/jit/mobile/import.h>
#include <torch/csrc/jit/mobile/module.h>
int main() {
std::string model_path = "AddTensorsModel.ptl";
torch::jit::ExtraFilesMap extra_files;
// Add keys into extra_files to indicate which metadata files
// we wish to fetch.
extra_files["model_info.txt"] = "";
extra_files["dtype_info.txt"] = "";
torch::jit::_load_extra_only_for_mobile(model_path, c10::nullopt, extra_files);
for (auto it: extra_files) {
std::cout << "Metadata [" << it.first << "]: " << it.second << std::endl;
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment