I would like to know if there is a better way to speed up the performance of my FFT over the first dimension of a tensor that is 1000 x 512 x 512, for example

I tried the following to overcome memory issues

device = 'cuda' if torch.cuda.is_available() else 'cpu'
z = torch.randn(1000,512,512)
for ii in tqdm(range(z.shape[-1]), total = z.shape[-1], desc='Applyinf FFT 1D'):
    for jj in range(z.shape[1]):
        z[...,jj, ii] = torch.fft.fftshift(torch.fft.fft(z[...,jj,ii].to(device), dim = 0), dim = 0).abs().cpu()

But it seems that the only thin for me to do is to batch in the last dimensions until I find an optimal compromise between time and memory usage for my fft