| 512 | |
| 513 | @pytest.mark.parametrize("test_traced_module", [True, False]) |
| 514 | def test_pickle_module(test_traced_module): |
| 515 | data_shape = (2, 28) |
| 516 | data = tensor(np.random.random(data_shape)) |
| 517 | mlp = MLP() |
| 518 | pred_gt = mlp(data) |
| 519 | if test_traced_module: |
| 520 | mlp = trace_module(mlp, data) |
| 521 | # pickle before forward |
| 522 | with BytesIO() as fout: |
| 523 | mge.save(mlp, fout) |
| 524 | fout.seek(0) |
| 525 | mlp1 = mge.load(fout) |
| 526 | if test_traced_module: |
| 527 | assert type(mlp1) == TracedModule |
| 528 | pred0 = mlp1(data) |
| 529 | |
| 530 | pred1 = mlp(data) |
| 531 | |
| 532 | # pickle after forward |
| 533 | with BytesIO() as fout: |
| 534 | mge.save(mlp, fout) |
| 535 | fout.seek(0) |
| 536 | mlp1 = mge.load(fout) |
| 537 | if test_traced_module: |
| 538 | assert type(mlp1) == TracedModule |
| 539 | pred2 = mlp1(data) |
| 540 | |
| 541 | np.testing.assert_allclose(pred_gt.numpy(), pred1.numpy(), atol=5e-6) |
| 542 | np.testing.assert_allclose(pred0.numpy(), pred1.numpy(), atol=5e-6) |
| 543 | np.testing.assert_allclose(pred0.numpy(), pred2.numpy(), atol=5e-6) |
| 544 | |
| 545 | |
| 546 | def test_repr_basic(): |