From c1e42565084bb1b110fb0c54ced3c86171e20e2e Mon Sep 17 00:00:00 2001 From: conti3000 Date: Thu, 25 Jan 2024 16:08:33 -0500 Subject: [PATCH] Update functional.py Green function and laplacion send to same device --- pietorch/functional.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pietorch/functional.py b/pietorch/functional.py index 7bd60cd..ed3d6eb 100644 --- a/pietorch/functional.py +++ b/pietorch/functional.py @@ -93,7 +93,8 @@ def blend(target: Tensor, source: Tensor, mask: Tensor, corner_coord: Tensor, mi else: assert green_function.shape[d] == 1, f'Green function should have size 1 in non-chosen dimension ' \ f'{d}: has {green_function.shape[d]}.' - + if laplacian.device != green_function.device: + green_function = green_function.to(laplacian.device) # Apply green function convolution init_blended = torch.fft.ifftn(torch.fft.fftn(laplacian, dim=chosen_dimensions) * green_function, dim=chosen_dimensions)