fix bugs of sum, add warning to topk

torch2trt_dynamic/converters/sum.py

@@ -1,24 +1,27 @@
 from torch2trt_dynamic.torch2trt_dynamic import *
 from torch2trt_dynamic.module_test import add_module_test
 from .unary import UnaryModule
-    
+
 
 @tensorrt_converter('torch.sum')
 @tensorrt_converter('torch.Tensor.sum')
 def convert_sum(ctx):
     input = ctx.method_args[0]
     dim = get_arg(ctx, 'dim', pos=1, default=tuple(range(1, input.ndim)))
     keepdim = get_arg(ctx, 'keepdim', pos=2, default=False)
-    input_trt= trt_(ctx.network, input)
+    if dim < 0:
+        dim = input.dim() + dim
+    input_trt = trt_(ctx.network, input)
     output = ctx.method_return
-    layer = ctx.network.add_reduce(input_trt,  trt.ReduceOperation.SUM, torch_dim_to_trt_axes(dim), keepdim)
+    layer = ctx.network.add_reduce(input_trt, trt.ReduceOperation.SUM,
+                                   torch_dim_to_trt_axes(dim), keepdim)
     output._trt = layer.get_output(0)
-        
-        
+
+
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 3)])
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 3, 3)])
 def test_sum_reduce_all():
-    return UnaryModule(lambda x: torch.sum(x))     
+    return UnaryModule(lambda x: torch.sum(x))
 
 
 @add_module_test(torch.float32, torch.device('cuda'), [(1, 3)])

torch2trt_dynamic/converters/topk.py

@@ -15,6 +15,12 @@ def convert_topk(ctx):
         axis = len(input.shape) - 1
     if axis < 0:
         axis = len(input.shape) + axis
+
+    if k > 3840:
+        print("warning: topk = " + k +
+              " > 3840 is not allowed in TensorRT, use 3840 instead.")
+        k = 3840
+
     largest = get_arg(ctx, 'largest', pos=3, default=True)
     topkOp = trt.TopKOperation.MAX if largest else trt.TopKOperation.MIN