xref: /aosp_15_r20/external/pytorch/torch/csrc/jit/serialization/import_read.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <torch/csrc/jit/serialization/unpickler.h>
4 #include <memory>
5 
6 namespace caffe2::serialize {
7 class PyTorchStreamReader;
8 } // namespace caffe2::serialize
9 
10 namespace torch::jit {
11 
12 TORCH_API IValue readArchiveAndTensors(
13     const std::string& archive_name,
14     const std::string& pickle_prefix,
15     const std::string& tensor_prefix,
16     std::optional<TypeResolver> type_resolver,
17     std::optional<ObjLoader> obj_loader,
18     std::optional<at::Device> device,
19     caffe2::serialize::PyTorchStreamReader& stream_reader,
20     c10::TypePtr (*type_parser)(const std::string&) =
21         Unpickler::defaultTypeParser,
22     std::shared_ptr<DeserializationStorageContext> storage_context = nullptr);
23 
24 bool check_zip_file(
25     const std::shared_ptr<caffe2::serialize::ReadAdapterInterface>& rai);
26 
27 } // namespace torch::jit
28