diff --git a/model_compression_toolkit/core/common/hessian/hessian_info_service.py b/model_compression_toolkit/core/common/hessian/hessian_info_service.py index 41bd896df..6d93dc126 100644 --- a/model_compression_toolkit/core/common/hessian/hessian_info_service.py +++ b/model_compression_toolkit/core/common/hessian/hessian_info_service.py @@ -204,7 +204,7 @@ def _fetch_hessians_with_compute(self, request: HessianScoresRequest, n_iteratio target_nodes = [n for n in orig_request.target_nodes if n.name in missing] request = request.clone(target_nodes=target_nodes) self._compute_hessians(request, n_iterations, count_by_cache=True) - res, missing = self.cache.fetch_hessian(request) + res, missing = self.cache.fetch_hessian(orig_request) assert not missing return res