I split my Fourier Neural Operator(FNO) model using model parallel(Single-Machine Model Parallel Best Practices — PyTorch Tutorials 2.4.0+cu124 documentation). The output from the model was used for calculating PDEs using Fourier Transform(torch.fft.fftn).

I was able to acquire the required PDE solution when I used only one GPU. However, the PDE solution increased significantly when the same code was run on two GPUs.

Is there a solution for this issue?

Compute the FFT of the model’s output, apply the differential operator in the frequency domain, and then perform an inverse FFT. Here’s an example:

```
import torch.fft
# Assuming output is your FNO model's output
fft_output = torch.fft.fftn(output)
# Apply the differential operator in the frequency domain
# Example for a second-order derivative: multiply by -k^2
k = torch.fft.fftfreq(output.size(-1))
k_squared = -k ** 2
fft_derivative = fft_output * k_squared
# Inverse FFT to get the derivative in spatial domain
derivative = torch.fft.ifftn(fft_derivative).real
```

The results I get from the above steps is different when I use 2 GPUs instead of 1 GPU.

Do you know why I am not getting the same result when I use 2 GPUs ?

@Gaijinliu_Gangmei It is quite possible that you are not setting the same manual_seed

```
torch.manual_seed(42)
```

The problem still exists even after using the same manual seed( torch.cuda.manual_seed_all(42) and torch.manual_seed(42)).