frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 05:34:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

[TorchScript] Support user defined classes as constants (#5062)

Browse Source 
Summary:
Pull Request resolved: https://github.com/pytorch/glow/pull/5062

Pull Request resolved: https://github.com/pytorch/pytorch/pull/45556

User defined classes can be used as constants.  This is useful when freezing and removing the module from the graph.

Test Plan: waitforsadcastle

Reviewed By: eellison

Differential Revision: D23994974

fbshipit-source-id: 5b4a5c91158aa7f22df39d71f2658afce1d29317
This commit is contained in:
Bram Wasti
2020-11-16 19:16:55 -08:00
committed by Facebook GitHub Bot
parent 3611d26a25
commit 43a9d6fb6e
17 changed files with 144 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
torch/quantization/quantize_jit.py
View File

@ -87,6 +87,7 @@ def _convert_jit(model, inplace=False, debug=False, quant_type=QuantType.STATIC,
model._reconstruct(model_c)
else:
model = wrap_cpp_module(model_c)
torch._C._jit_pass_constant_propagation(model.graph)
return model
def convert_jit(model, inplace=False, debug=False, preserved_attrs=None):
@ -110,6 +111,7 @@ def _quantize_jit(model, qconfig_dict, run_fn=None, run_args=None, inplace=False
run_fn(model, *run_args)
model = convert_jit(model, True, debug)
torch._C._jit_pass_constant_propagation(model.graph)
return model
def quantize_jit(model, qconfig_dict, run_fn, run_args, inplace=False, debug=False):