update

Signed-off-by: Wenqi Li <wenqil@nvidia.com>

Author: wyli

monai/networks/utils.py

@@ -854,7 +854,6 @@ class AutoDeviceAdapt:
     calling ``mod(inputs, ...)`` with heterogeneous ``inputs`` on CUDA.
     Since the error handling is expensive, if the OOM happens frequently, ``mod`` should be adjusted instead of
     relying on this class.
-
     """
 
     def __init__(self, mod):
@@ -890,13 +889,16 @@ def __call__(self, inputs, *args, **kwargs):
                 f" please review {type(self.mod)} settings if this message appears frequently."
             )
             if not self.has_cpu_thresh:  # mod doesn't support dynamic input size based device allocation
+                if isinstance(inputs, torch.Tensor):
+                    inputs = inputs.to("cpu")
                 if self.has_device:  # run this call with mod.device="cpu"
                     ori_device = self.mod.device
                     self.mod.device = "cpu"
-                    out = self.mod(inputs.to("cpu"), *args, **kwargs)
+                    out = self.mod(inputs, *args, **kwargs)
                     self.mod.device = ori_device
                     return out
-                if isinstance(inputs, torch.Tensor):
-                    return self.mod(inputs.to("cpu"), *args, **kwargs)  # move inputs to cpu
-            self.mod.cpu_thresh = inputs.shape[2:].numel() - 1  # try change the dynamic threshold before the call
-            return self.mod(inputs, *args, **kwargs)
+                return self.mod(inputs, *args, **kwargs)
+            if isinstance(inputs, torch.Tensor):
+                self.mod.cpu_thresh = inputs.shape[2:].numel() - 1  # try change the dynamic threshold before the call
+                return self.mod(inputs, *args, **kwargs)
+            raise e