From efc05f2b0552bb6beb38641e1112b29c94b37eba Mon Sep 17 00:00:00 2001 From: Alexander Suslov Date: Thu, 1 Feb 2024 14:52:48 +0400 Subject: [PATCH] Removed model state recovering after model tracing (#2419) ### Changes Removed model state recovering after model tracing ### Reason for changes Reduce memory usage ### Related tickets N/A ### Tests e2e_pytorch_full 578 --- nncf/torch/dynamic_graph/graph_tracer.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/nncf/torch/dynamic_graph/graph_tracer.py b/nncf/torch/dynamic_graph/graph_tracer.py index 505f3dd8440..4a86cb68576 100644 --- a/nncf/torch/dynamic_graph/graph_tracer.py +++ b/nncf/torch/dynamic_graph/graph_tracer.py @@ -8,7 +8,6 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from copy import deepcopy from typing import Any, Callable, Dict, Optional, Tuple, TypeVar import torch @@ -35,8 +34,6 @@ def trace_graph( as_eval: bool = False, trace_parameters: bool = False, ) -> DynamicGraph: - sd = deepcopy(model.state_dict()) - if context_to_use is None: context_to_use = TracingContext() @@ -54,7 +51,6 @@ def trace_graph( self.custom_forward_fn(model) else: self.custom_forward_fn(model) - model.load_state_dict(sd) context_to_use.disable_trace_dynamic_graph() return context_to_use.graph