| 620 | self.assertTrue(torch.allclose(ep(*inp), loaded_ep(*inp))) |
| 621 | |
| 622 | def test_save_extra(self): |
| 623 | inp = (torch.tensor([0.1, 0.1]),) |
| 624 | |
| 625 | def f(x): |
| 626 | return x * x + x |
| 627 | |
| 628 | ep = export(f, inp) |
| 629 | |
| 630 | buffer = io.BytesIO() |
| 631 | save(ep, buffer, extra_files={"extra.txt": "moo"}) |
| 632 | buffer.seek(0) |
| 633 | extra_files = {"extra.txt": ""} |
| 634 | loaded_ep = load(buffer, extra_files=extra_files) |
| 635 | |
| 636 | self.assertTrue(torch.allclose(ep(*inp), loaded_ep(*inp))) |
| 637 | self.assertEqual(extra_files["extra.txt"], "moo") |
| 638 | |
| 639 | def test_version_error(self): |
| 640 | def f(x): |