mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
a4a6d056e6
Summary: Fixes https://github.com/pytorch/pytorch/issues/66232 This should be the last immediate task. I anticipate test ownership will change overtime but this is the last big thing to close it out Pull Request resolved: https://github.com/pytorch/pytorch/pull/67859 Reviewed By: soulitzer Differential Revision: D32210534 Pulled By: janeyx99 fbshipit-source-id: 7fd835d87d9d35d49ec49de1fcfa29b085133e99
37 lines
1.0 KiB
Python
37 lines
1.0 KiB
Python
# Owner(s): ["oncall: mobile"]
|
|
|
|
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()
|