Merge pull request #397 from NVIDIA-AI-IOT/const_dim_fix

Const dim fix

Author: jaybdub

torch2trt/converters/AdaptiveAvgPool2d.py

@@ -8,7 +8,7 @@ def convert_AdaptiveAvgPool2d(ctx):
     input = ctx.method_args[1]
     output = ctx.method_return
 
-    input_trt = trt_(ctx.network, input)
+    input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
 
     output_size = module.output_size
     if not isinstance(output_size, tuple):

torch2trt/converters/BatchNorm1d.py

@@ -6,7 +6,7 @@
 def convert_BatchNorm2d(ctx):
     module = ctx.method_args[0]
     input = ctx.method_args[1]
-    input_trt = trt_(ctx.network, input)
+    input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
     output = ctx.method_return
 
     scale = module.weight.detach().cpu().numpy() / np.sqrt(module.running_var.detach().cpu().numpy() + module.eps)

torch2trt/converters/BatchNorm2d.py

@@ -6,7 +6,7 @@
 def convert_BatchNorm2d(ctx):
     module = ctx.method_args[0]
     input = ctx.method_args[1]
-    input_trt = trt_(ctx.network, input)
+    input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
     output = ctx.method_return
 
     scale = module.weight.detach().cpu().numpy() / np.sqrt(

torch2trt/converters/Conv.py

@@ -7,7 +7,7 @@
 def convert_Conv_trt7(ctx):
     module = ctx.method_args[0]
     input = ctx.method_args[1]
-    input_trt = trt_(ctx.network, input)
+    input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
     output = ctx.method_return
 
     input_dim = input.dim() - 2

torch2trt/converters/Conv1d.py

@@ -6,7 +6,7 @@
 def convert_Conv1d(ctx):
     module = ctx.method_args[0]
     input = ctx.method_args[1]
-    input_trt = trt_(ctx.network, input)
+    input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
     output = ctx.method_return
 
     kernel_size = (module.kernel_size[0], 1)

torch2trt/converters/Conv2d.py

@@ -6,7 +6,7 @@
 def convert_Conv2d(ctx):
     module = ctx.method_args[0]
     input = ctx.method_args[1]
-    input_trt = trt_(ctx.network, input)
+    input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
     output = ctx.method_return
 
     kernel_size = module.kernel_size

torch2trt/converters/ConvTranspose.py

@@ -7,7 +7,7 @@
 def convert_ConvTranspose2d_trt7(ctx):
     module = ctx.method_args[0]
     input = ctx.method_args[1]
-    input_trt = trt_(ctx.network, input)
+    input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
     output = ctx.method_return
 
     input_dim = input.dim() - 2

torch2trt/converters/ConvTranspose2d.py

@@ -5,7 +5,7 @@
 def convert_ConvTranspose2d(ctx):
     module = ctx.method_args[0]
     input = ctx.method_args[1]
-    input_trt = trt_(ctx.network, input)
+    input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
     output = ctx.method_return
 
     kernel_size = module.kernel_size

torch2trt/converters/Identity.py

@@ -6,6 +6,6 @@
 @tensorrt_converter('torch.nn.Dropout3d.forward')
 def convert_Identity(ctx):
     input = ctx.method_args[1]
-    input_trt = trt_(ctx.network, input)
+    input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
     output = ctx.method_return
     output._trt = input_trt

torch2trt/converters/Linear.py

@@ -6,7 +6,7 @@
 def convert_Linear(ctx):
     module = ctx.method_args[0]
     input = ctx.method_args[1]
-    input_trt = trt_(ctx.network, input)
+    input_trt = add_missing_trt_tensors(ctx.network, [input])[0]
     output = ctx.method_return
 
     # reshape to ...xNx1x1