| 570 | @unittest.skipIf(not torchdynamo.is_dynamo_supported(), "dynamo doesn't support") |
| 571 | class TestSaveLoad(TestCase): |
| 572 | def test_save_buffer(self): |
| 573 | inp = (torch.tensor([0.1, 0.1]),) |
| 574 | linear = torch.nn.Linear(2, 2) |
| 575 | |
| 576 | class Module(torch.nn.Module): |
| 577 | def forward(self, x): |
| 578 | x = x + 1 |
| 579 | y = x.t() |
| 580 | y = y.relu() |
| 581 | y = linear(y) |
| 582 | return y |
| 583 | |
| 584 | ep = export(Module(), inp) |
| 585 | |
| 586 | buffer = io.BytesIO() |
| 587 | save(ep, buffer) |
| 588 | buffer.seek(0) |
| 589 | loaded_ep = load(buffer) |
| 590 | |
| 591 | self.assertTrue(torch.allclose(ep(*inp), loaded_ep(*inp))) |
| 592 | |
| 593 | def test_save_file(self): |
| 594 | |