From f507675094ff82bb303b5393023ff3711df88a2b Mon Sep 17 00:00:00 2001 From: ytl0623 Date: Tue, 10 Feb 2026 13:48:53 +0800 Subject: [PATCH 1/2] Fix AutoencoderKlMaisi forcing CUDA transfer on CPU inputs Signed-off-by: ytl0623 --- monai/apps/generation/maisi/networks/autoencoderkl_maisi.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py b/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py index 16697204c5..64e0fc046f 100644 --- a/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py +++ b/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py @@ -214,6 +214,8 @@ def _concatenate_tensors(self, outputs: list[torch.Tensor], split_size: int, pad if max(outputs[0].size()) < 500: x = torch.cat(outputs, dim=self.dim_split + 2) else: + target_device = outputs[0].device + x = outputs[0].clone().to("cpu", non_blocking=True) outputs[0] = torch.Tensor(0) _empty_cuda_cache(self.save_mem) @@ -225,7 +227,9 @@ def _concatenate_tensors(self, outputs: list[torch.Tensor], split_size: int, pad if self.print_info: logger.info(f"MaisiConvolution concat progress: {k + 1}/{len(outputs) - 1}.") - x = x.to("cuda", non_blocking=True) + if target_device.type != "cpu": + x = x.to(target_device, non_blocking=True) + return x def forward(self, x: torch.Tensor) -> torch.Tensor: From a70d48d602a9b699bb86c23af393472c9e0dc794 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 10 Feb 2026 05:51:41 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- monai/apps/generation/maisi/networks/autoencoderkl_maisi.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py b/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py index 64e0fc046f..39f8459224 100644 --- a/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py +++ b/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py @@ -215,7 +215,7 @@ def _concatenate_tensors(self, outputs: list[torch.Tensor], split_size: int, pad x = torch.cat(outputs, dim=self.dim_split + 2) else: target_device = outputs[0].device - + x = outputs[0].clone().to("cpu", non_blocking=True) outputs[0] = torch.Tensor(0) _empty_cuda_cache(self.save_mem) @@ -229,7 +229,7 @@ def _concatenate_tensors(self, outputs: list[torch.Tensor], split_size: int, pad if target_device.type != "cpu": x = x.to(target_device, non_blocking=True) - + return x def forward(self, x: torch.Tensor) -> torch.Tensor: