mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
ba48cf6535
See https://github.com/pytorch/pytorch/pull/129751#issue-2380881501. Most changes are auto-generated by linter. You can review these PRs via: ```bash git diff --ignore-all-space --ignore-blank-lines HEAD~1 ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/129757 Approved by: https://github.com/ezyang
39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
# Owner(s): ["oncall: mobile"]
|
|
|
|
import io
|
|
import tempfile
|
|
import unittest
|
|
|
|
import torch
|
|
import torch.utils.show_pickle
|
|
from torch.testing._internal.common_utils import IS_WINDOWS, run_tests, TestCase
|
|
|
|
|
|
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()
|