handle 1d case batch_norm

Author: jaybdub

torch2trt/converters/native_converters.py

@@ -148,14 +148,38 @@ def convert_batch_norm(ctx):
     bias = get_arg(ctx, 'bias', pos=4, default=None) 
     eps = get_arg(ctx, 'eps', pos=7, default=10e-6) 
 
+    ndim = input.ndim - 2
     input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
     output = ctx.method_return
 
     scale = weight.detach().cpu().numpy() / np.sqrt(running_var.detach().cpu().numpy() + eps)
     bias = bias.detach().cpu().numpy() - running_mean.detach().cpu().numpy() * scale
     power = np.ones_like(scale)
 
-    layer = ctx.network.add_scale_nd(input_trt, trt.ScaleMode.CHANNEL, bias, scale, power, 1)
+    if ndim == 1:
+        # reshape to 2D
+        layer = ctx.network.add_shuffle(input_trt)
+        
+        if len(input.shape) == 2:
+            layer.reshape_dims = (0, 0, 1, 1)
+        else:
+            layer.reshape_dims = (0, 0, 0, 1)
+
+        scale_input = layer.get_output(0)
+    else:
+        scale_input = input_trt
+
+    layer = ctx.network.add_scale_nd(scale_input, trt.ScaleMode.CHANNEL, bias, scale, power, 1)
+
+    if ndim == 1:
+        # reshape back to 1D
+        layer = ctx.network.add_shuffle(layer.get_output(0))
+        if len(input.shape) == 2:
+            layer.reshape_dims = (0, 0)
+        else:
+            layer.reshape_dims = (0, 0, 0)
+
+
     output._trt = layer.get_output(0)