mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
fc5db4265b
Summary: fix https://github.com/pytorch/pytorch/issues/50448. This replaces all `test/*.py` files with run_tests(). This PR does not address test files in the subdirectories because they seems unrelated. Pull Request resolved: https://github.com/pytorch/pytorch/pull/50451 Reviewed By: janeyx99 Differential Revision: D25899924 Pulled By: walterddr fbshipit-source-id: f7c861f0096624b2791ad6ef6a16b1c4895cce71
35 lines
1020 B
Python
35 lines
1020 B
Python
import unittest
|
|
import io
|
|
import tempfile
|
|
import torch
|
|
import torch.utils.show_pickle
|
|
|
|
from torch.testing._internal.common_utils import TestCase, run_tests, IS_WINDOWS
|
|
|
|
class TestShowPickle(TestCase):
|
|
|
|
@unittest.skipIf(IS_WINDOWS, "Can't re-open temp file on Windows")
|
|
def test_scripted_model(self):
|
|
class MyCoolModule(torch.nn.Module):
|
|
def __init__(self, weight):
|
|
super().__init__()
|
|
self.weight = weight
|
|
|
|
def forward(self, x):
|
|
return x * self.weight
|
|
|
|
m = torch.jit.script(MyCoolModule(torch.tensor([2.0])))
|
|
|
|
with tempfile.NamedTemporaryFile() as tmp:
|
|
torch.jit.save(m, tmp)
|
|
tmp.flush()
|
|
buf = io.StringIO()
|
|
torch.utils.show_pickle.main(["", tmp.name + "@*/data.pkl"], output_stream=buf)
|
|
output = buf.getvalue()
|
|
self.assertRegex(output, "MyCoolModule")
|
|
self.assertRegex(output, "weight")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
run_tests()
|