xref: /aosp_15_r20/external/pytorch/test/load_torchscript_model.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1import sys
2
3import torch
4
5
6if __name__ == "__main__":
7    script_mod = torch.jit.load(sys.argv[1])
8    mod = torch.load(sys.argv[1] + ".orig")
9    print(script_mod)
10    inp = torch.rand(2, 28 * 28)
11    _ = mod(inp)
12    sys.exit(0)
13