From 3152a271d375892e06ba3bfe3286152061a6f6ce Mon Sep 17 00:00:00 2001 From: Takeshi Watanabe Date: Tue, 19 Nov 2024 11:37:57 +0900 Subject: [PATCH] Transfer to cpu only when torch.Tensor arguments in onnx export --- mmdeploy/apis/onnx/export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mmdeploy/apis/onnx/export.py b/mmdeploy/apis/onnx/export.py index 76c4fb1003..790d002509 100644 --- a/mmdeploy/apis/onnx/export.py +++ b/mmdeploy/apis/onnx/export.py @@ -132,7 +132,7 @@ def wrapper(*arg, **kwargs): if isinstance(args, torch.Tensor): args = args.cpu() elif isinstance(args, (tuple, list)): - args = tuple([_.cpu() for _ in args]) + args = tuple([_.cpu() if isinstance(_, torch.Tensor) else _ for _ in args]) else: raise RuntimeError(f'Not supported args: {args}') torch.onnx.export(