Fix pytorch upsample parsing (fastmachinelearning#1186)

* Fix parsing Pytorch Upsample layer in 2D

* Add test case for Upsample to have a layer after to ensure proper parsing

Author: vloncar

hls4ml/converters/pytorch/reshape.py

@@ -152,7 +152,7 @@ def handle_upsample(operation, layer_name, input_names, input_shapes, node, clas
         layer['out_height'] = int(layer['in_height'] * scale_height)
         layer['out_width'] = int(layer['in_width'] * scale_width)
 
-        output_shape = [layer['n_chan'], layer['out_height'], layer['out_width']]
+        output_shape = [input_shapes[0][0], layer['n_chan'], layer['out_height'], layer['out_width']]
     else:
         raise Exception(f'Parsing "Upsample" with {len(input_shape)}-dimensional tensors is not yet supported.')
 

test/pytest/test_upsampling_pytorch.py

@@ -33,19 +33,21 @@ class Upsample1DModel(nn.Module):
     def __init__(self):
         super().__init__()
         self.upsample = nn.Upsample(scale_factor=2)
+        self.relu = nn.ReLU()
 
     def forward(self, x):
-        return self.upsample(x)
+        return self.relu(self.upsample(x))
 
 
 class Upsample2DModel(nn.Module):
     def __init__(self):
         super().__init__()
         # this scale_factor tests proper output shape calculation with fractional scaling and parsing per-axis scales
         self.upsample = nn.UpsamplingNearest2d(scale_factor=(1, 2.4))  # Would also work with Upsample(mode='nearest')
+        self.relu = nn.ReLU()
 
     def forward(self, x):
-        return self.upsample(x)
+        return self.relu(self.upsample(x))
 
 
 @pytest.mark.parametrize('io_type', ['io_stream', 'io_parallel'])