Skip to content

Commit

Permalink
Add cleanup_torchscript_cache to test_quantize_conformance.py (#2171)
Browse files Browse the repository at this point in the history
### Changes

Add `cleanup_torchscript_cache` function to test_quantize_conformance.py

### Reason for changes

After run torch.jit.trace in convert_model, PyTorch does not clear the
trace cache automatically.

Same function in torch test and openvino notebok:

https://github.com/pytorch/pytorch/blob/main/torch/testing/_internal/jit_utils.py#L59

https://github.com/openvinotoolkit/openvino_notebooks/blob/1932d4b4e99116bdedaa620c9dc92069fbb1f05e/notebooks/236-stable-diffusion-v2/implementation/conversion_helper_utils.py#L8
  • Loading branch information
AlexanderDokuchaev authored Oct 9, 2023
1 parent d39d3ae commit a97fff9
Showing 1 changed file with 13 additions and 0 deletions.
13 changes: 13 additions & 0 deletions tests/post_training/pipelines/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -279,6 +279,19 @@ def run(self) -> None:
self.save_quantized_model()
self.get_num_fq()
self.validate()
self.cleanup_torchscript_cache()

@staticmethod
def cleanup_torchscript_cache():
"""
Helper for removing cached model representation.
After run torch.jit.trace in convert_model, PyTorch does not clear the trace cache automatically.
"""
# pylint: disable=protected-access
torch._C._jit_clear_class_registry()
torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore()
torch.jit._state._clear_class_state()

def get_run_info(self) -> RunInfo:
return self.run_info

0 comments on commit a97fff9

Please sign in to comment.