fix bugs of zeros ,ones

grimoire · Sep 26, 2020 · 1f247bf · 1f247bf
1 parent 3702fbd
commit 1f247bf
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 0 deletions.
diff --git a/torch2trt/converters/new_ones.py b/torch2trt/converters/new_ones.py
@@ -9,6 +9,9 @@ def convert_new_ones(ctx):
 
     output = ctx.method_return
 
+    if isinstance(size, int):
+        size = (size, )
+
     # check const
     is_const = True
     for s in size:

diff --git a/torch2trt/converters/new_zeros.py b/torch2trt/converters/new_zeros.py
@@ -9,6 +9,9 @@ def convert_new_zeros(ctx):
 
     output = ctx.method_return
 
+    if isinstance(size, int):
+        size = (size, )
+
     # check const
     is_const = True
     for s in size:

diff --git a/torch2trt/converters/zeros.py b/torch2trt/converters/zeros.py
@@ -11,6 +11,9 @@ def convert_zeros(ctx):
         dtype = ctx.method_kwargs["dtype"]
     output = ctx.method_return
 
+    if isinstance(size, int):
+        size = (size, )
+
     # check const
     is_const = True
     for s in size: