ensure inputs to graphs are float in onnx and tensorrt export

Clarifai · Dec 30, 2021 · f62cf69 · f62cf69
1 parent 7e24fa7
commit f62cf69
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/tools/onnx2tensorrt.py b/tools/onnx2tensorrt.py
@@ -40,7 +40,7 @@ def _prepare_input_img(img_path: str,
     # prepare data
     data = dict(img=img_path)
     data = test_pipeline(data)
-    imgs = data['img']
+    imgs = [img.float() for img in data['img']]
     img_metas = [i.data for i in data['img_metas']]
 
     if rescale_shape is not None:

diff --git a/tools/pytorch2onnx.py b/tools/pytorch2onnx.py
@@ -90,7 +90,7 @@ def _prepare_input_img(img_path,
     # prepare data
     data = dict(img=img_path)
     data = test_pipeline(data)
-    imgs = data['img']
+    imgs = [img.float() for img in data['img']]
     img_metas = [i.data for i in data['img_metas']]
 
     if rescale_shape is not None: