(self)
| 591 | self.assertTrue(torch.allclose(ep(*inp), loaded_ep(*inp))) |
| 592 | |
| 593 | def test_save_file(self): |
| 594 | |
| 595 | def f(x): |
| 596 | return x * x |
| 597 | |
| 598 | inp = (torch.randn(2, 2),) |
| 599 | ep = export(f, inp) |
| 600 | |
| 601 | with tempfile.NamedTemporaryFile() as f: |
| 602 | save(ep, f) |
| 603 | f.seek(0) |
| 604 | loaded_ep = load(f) |
| 605 | |
| 606 | self.assertTrue(torch.allclose(ep(*inp), loaded_ep(*inp))) |
| 607 | |
| 608 | def test_save_path(self): |
| 609 | def f(x, y): |