# How to calculate PDEs(via torch.fft.fftn) from the output of FNO model using model parallel?

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)).