| import torch | |
| from torch.testing._internal.jit_utils import JitTestCase | |
| if __name__ == '__main__': | |
| raise RuntimeError("This test file is not meant to be run directly, use:\n\n" | |
| "\tpython test/test_jit.py TESTNAME\n\n" | |
| "instead.") | |
| class TestPythonIr(JitTestCase): | |
| def test_param_strides(self): | |
| def trace_me(arg): | |
| return arg | |
| t = torch.zeros(1, 3, 16, 16) | |
| traced = torch.jit.trace(trace_me, t) | |
| value = list(traced.graph.param_node().outputs())[0] | |
| real_strides = list(t.stride()) | |
| type_strides = value.type().strides() | |
| self.assertEqual(real_strides, type_strides) |