Merge pull request #56 from TCLResearchEurope/fix-requires-grad

Fix requires grad

Author: michal-lopuszynski-tcl

torch_dag/core/unstructured_to_structured.py

@@ -236,7 +236,7 @@ def build_dag_by_tracing(self) -> dag_module.DagModule:
         for node in graph.nodes:
             num_predecessors = len(node._input_nodes)
             if num_predecessors == 0 and node.target in model.state_dict():  # free parameter node
-                module = structured_modules.ParameterModule(param=torch.nn.Parameter(model.state_dict()[node.target]))
+                module = structured_modules.ParameterModule(param=torch.nn.Parameter(model.get_parameter(node.target), requires_grad = model.get_parameter(node.target).requires_grad))
                 vertex = dag.add_vertex(
                     name=node.name,
                     module=module,
@@ -395,6 +395,9 @@ def convert_node(self, node: torch.fx.node.Node, modules_dict, state_dict):
         elif node.target == torch.square:
             return structured_modules.PowerModule(pow=2)
 
+        elif node.target == torch.add:
+            return structured_modules.AddModule()
+
         elif node.target == torch.addcmul:
             return structured_modules.AddcmulModule()