fix missing input names in squeeze layers

fastmachinelearning · Dec 13, 2024 · 3e5b7eb · 3e5b7eb
1 parent f58ed42
commit 3e5b7eb
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/hls4ml/converters/pytorch/reshape.py b/hls4ml/converters/pytorch/reshape.py
@@ -38,6 +38,7 @@ def parse_squeeze_layer(operation, layer_name, input_names, input_shapes, node,
     layer = {}
     layer['class_name'] = 'Reshape'
     layer['name'] = layer_name
+    layer['inputs'] = input_names
 
     if len(node.args) > 1 or len(node.kwargs) > 0:  # 'dim' argument is specified
         output_shape = [i for i in input_shapes[0]]