leap-model-parser 0.1.250.dev6__tar.gz → 0.1.250.dev7__tar.gz

{leap_model_parser-0.1.250.dev6 → leap_model_parser-0.1.250.dev7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: leap-model-parser
-Version: 0.1.250.dev6
+Version: 0.1.250.dev7
 Summary: 
 Home-page: https://github.com/tensorleap/leap-model-parser
 License: MIT

{leap_model_parser-0.1.250.dev6 → leap_model_parser-0.1.250.dev7}/leap_model_parser/torch_utils.py

@@ -55,17 +55,28 @@ def extract_tensor(output: Any) -> torch.Tensor:
 
 
 class ActivationInterpreter(fx.Interpreter):
-    def __init__(self, gm: fx.GraphModule):
+    def __init__(self, gm: fx.GraphModule, grad_mode: bool = False):
         super().__init__(gm)
         self.cache: Dict[str, torch.Tensor] = {}
+        self._grad_mode = grad_mode
 
     def _store_tensor(self, name: str, value: Any) -> None:
         if isinstance(value, torch.Tensor):
-            self.cache[name] = value.detach().clone()
+            if self._grad_mode:
+                self.cache[name] = value
+                if value.requires_grad:
+                    value.retain_grad()
+            else:
+                self.cache[name] = value.detach().clone()
         elif isinstance(value, (tuple, list)):
             for item in value:
                 if isinstance(item, torch.Tensor):
-                    self.cache[name] = item.detach().clone()
+                    if self._grad_mode:
+                        self.cache[name] = item
+                        if item.requires_grad:
+                            item.retain_grad()
+                    else:
+                        self.cache[name] = item.detach().clone()
                     break
 
     def run_node(self, node: fx.Node) -> Any:

{leap_model_parser-0.1.250.dev6 → leap_model_parser-0.1.250.dev7}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "leap-model-parser"
-version = "0.1.250.dev6"
+version = "0.1.250.dev7"
 description = ""
 authors = ["idan <idan.yogev@tensorleap.ai>"]
 license = "MIT"