(self, bundle_files, bundle_name, version, repo, device, model_file)
| 401 | @parameterized.expand([TEST_CASE_9]) |
| 402 | @skip_if_quick |
| 403 | def test_load_ts_module(self, bundle_files, bundle_name, version, repo, device, model_file): |
| 404 | with skip_if_downloading_fails(): |
| 405 | # load ts module |
| 406 | with tempfile.TemporaryDirectory() as tempdir: |
| 407 | bundle_root = os.path.join(tempdir, bundle_name) |
| 408 | # load ts module |
| 409 | model_ts, metadata, extra_file_dict = load( |
| 410 | name=bundle_name, |
| 411 | version=version, |
| 412 | model_file=model_file, |
| 413 | load_ts_module=True, |
| 414 | bundle_dir=tempdir, |
| 415 | repo=repo, |
| 416 | progress=False, |
| 417 | device=device, |
| 418 | source="github", |
| 419 | config_files=("network.json",), |
| 420 | ) |
| 421 | |
| 422 | # prepare and test ts |
| 423 | input_tensor = torch.load( |
| 424 | os.path.join(bundle_root, bundle_files[1]), map_location=device, weights_only=True |
| 425 | ) |
| 426 | output = model_ts.forward(input_tensor) |
| 427 | expected_output = torch.load( |
| 428 | os.path.join(bundle_root, bundle_files[0]), map_location=device, weights_only=True |
| 429 | ) |
| 430 | assert_allclose(output, expected_output, atol=1e-3, rtol=1e-3, type_test=False) |
| 431 | # test metadata |
| 432 | self.assertTrue(metadata["pytorch_version"] == "1.7.1") |
| 433 | # test extra_file_dict |
| 434 | self.assertTrue("network.json" in extra_file_dict.keys()) |
| 435 | |
| 436 | |
| 437 | class TestDownloadLargefiles(unittest.TestCase): |
nothing calls this directly
no test coverage detected