mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
20be31de90
Test Plan: revert-hammer Differential Revision: D35423079 (d4a709be3d) Original commit changeset: 31abc4364d5f Original Phabricator Diff: D35423079 (d4a709be3d) fbshipit-source-id: 09ca72ebc330088fbfdcf2cabce3b6385c948a47 (cherry picked from commit d458172fb33135243e5cb1a04a5bee9f7d242f30)
51 lines
1.3 KiB
Python
51 lines
1.3 KiB
Python
# Owner(s): ["oncall: package/deploy"]
|
|
|
|
from io import BytesIO
|
|
|
|
from torch.package import (
|
|
PackageExporter,
|
|
PackageImporter,
|
|
sys_importer,
|
|
)
|
|
from torch.testing._internal.common_utils import run_tests
|
|
|
|
try:
|
|
from .common import PackageTestCase
|
|
except ImportError:
|
|
# Support the case where we run this file directly.
|
|
from common import PackageTestCase
|
|
|
|
|
|
class TestRepackage(PackageTestCase):
|
|
"""Tests for repackaging."""
|
|
|
|
def test_repackage_import_indirectly_via_parent_module(self):
|
|
from package_d.imports_directly import ImportsDirectlyFromSubSubPackage
|
|
from package_d.imports_indirectly import ImportsIndirectlyFromSubPackage
|
|
|
|
model_a = ImportsDirectlyFromSubSubPackage()
|
|
buffer = BytesIO()
|
|
with PackageExporter(buffer) as pe:
|
|
pe.intern("**")
|
|
pe.save_pickle("default", "model.py", model_a)
|
|
|
|
buffer.seek(0)
|
|
pi = PackageImporter(buffer)
|
|
loaded_model = pi.load_pickle("default", "model.py")
|
|
|
|
model_b = ImportsIndirectlyFromSubPackage()
|
|
buffer = BytesIO()
|
|
with PackageExporter(
|
|
buffer,
|
|
importer=(
|
|
pi,
|
|
sys_importer,
|
|
),
|
|
) as pe:
|
|
pe.intern("**")
|
|
pe.save_pickle("default", "model_b.py", model_b)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_tests()
|