1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
|
# Owner(s): ["oncall: package/deploy"]
from pathlib import Path
from unittest import skipIf
from torch.package import PackageImporter
from torch.testing._internal.common_utils import IS_FBCODE, IS_SANDCASTLE, run_tests
try:
from .common import PackageTestCase
except ImportError:
# Support the case where we run this file directly.
from common import PackageTestCase
packaging_directory = f"{Path(__file__).parent}/package_bc"
class TestLoadBCPackages(PackageTestCase):
"""Tests for checking loading has backwards compatiblity"""
@skipIf(
IS_FBCODE or IS_SANDCASTLE,
"Tests that use temporary files are disabled in fbcode",
)
def test_load_bc_packages_nn_module(self):
"""Tests for backwards compatible nn module"""
importer1 = PackageImporter(f"{packaging_directory}/test_nn_module.pt")
loaded1 = importer1.load_pickle("nn_module", "nn_module.pkl")
@skipIf(
IS_FBCODE or IS_SANDCASTLE,
"Tests that use temporary files are disabled in fbcode",
)
def test_load_bc_packages_torchscript_module(self):
"""Tests for backwards compatible torchscript module"""
importer2 = PackageImporter(f"{packaging_directory}/test_torchscript_module.pt")
loaded2 = importer2.load_pickle("torchscript_module", "torchscript_module.pkl")
@skipIf(
IS_FBCODE or IS_SANDCASTLE,
"Tests that use temporary files are disabled in fbcode",
)
def test_load_bc_packages_fx_module(self):
"""Tests for backwards compatible fx module"""
importer3 = PackageImporter(f"{packaging_directory}/test_fx_module.pt")
loaded3 = importer3.load_pickle("fx_module", "fx_module.pkl")
if __name__ == "__main__":
run_tests()
|
