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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: leap-model-parser
-Version: 0.1.250.dev6
+Version: 0.1.250.dev8
 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.dev8}/leap_model_parser/torch_utils.py

@@ -26,10 +26,13 @@ def get_full_args(ep: ExportedProgram, user_inputs: List[torch.Tensor]) -> List[
 
 
 def build_module_to_node_map(ep: ExportedProgram) -> Dict[str, Tuple[str, str]]:
+    import operator as _op
     mapping: Dict[str, Tuple[str, str]] = {}
     for node in ep.graph_module.graph.nodes:
         if node.op != "call_function":
             continue
+        if node.target is _op.getitem:
+            continue
         stack = node.meta.get("nn_module_stack", {})
         if not stack:
             continue
@@ -55,17 +58,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.dev8}/pyproject.toml

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