mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-15 21:00:47 +00:00
046e88a291
Rewrite Python built-in class `super()` calls. Only non-semantic changes should be applied. - #94587 - #94588 - #94592 Also, methods with only a `super()` call are removed: ```diff class MyModule(nn.Module): - def __init__(self): - super().__init__() - def forward(self, ...): ... ``` Some cases that change the semantics should be kept unchanged. E.g.:f152a79be9/caffe2/python/net_printer.py (L184-L190)f152a79be9/test/test_jit_fuser_te.py (L2628-L2635)Pull Request resolved: https://github.com/pytorch/pytorch/pull/94592 Approved by: https://github.com/ezyang, https://github.com/seemethere
44 lines
1.1 KiB
Python
44 lines
1.1 KiB
Python
import argparse
|
|
import os.path
|
|
import sys
|
|
|
|
import torch
|
|
|
|
|
|
def get_custom_op_library_path():
|
|
if sys.platform.startswith("win32"):
|
|
library_filename = "custom_ops.dll"
|
|
elif sys.platform.startswith("darwin"):
|
|
library_filename = "libcustom_ops.dylib"
|
|
else:
|
|
library_filename = "libcustom_ops.so"
|
|
path = os.path.abspath("build/{}".format(library_filename))
|
|
assert os.path.exists(path), path
|
|
return path
|
|
|
|
|
|
class Model(torch.jit.ScriptModule):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.p = torch.nn.Parameter(torch.eye(5))
|
|
|
|
@torch.jit.script_method
|
|
def forward(self, input):
|
|
return torch.ops.custom.op_with_defaults(input)[0] + 1
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Serialize a script module with custom ops"
|
|
)
|
|
parser.add_argument("--export-script-module-to", required=True)
|
|
options = parser.parse_args()
|
|
|
|
torch.ops.load_library(get_custom_op_library_path())
|
|
|
|
model = Model()
|
|
model.save(options.export_script_module_to)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|