Overview

PyTorch 1.8, released on March 5, 2021, introduces the torch.fft module for Fourier transforms and adds official AMD GPU support via ROCm.

Main Features

torch.fft module

The new torch.fft module provides fast Fourier transform functions natively integrated into PyTorch, with GPU computation and autograd support for backpropagation.

python 
import torch

# Sinusoidal signal
t = torch.linspace(0, 1, 1000)
signal = torch.sin(2 * torch.pi * 50 * t)

# Fourier transform
spectrum = torch.fft.fft(signal)
frequencies = torch.fft.fftfreq(len(t), d=1/1000)
print(f'Dominant frequency: {frequencies[spectrum.abs().argmax()]:.0f} Hz')
# Dominant frequency: 50 Hz

AMD ROCm support

PyTorch 1.8 adds official support for the AMD ROCm platform, allowing the same models to run on AMD GPUs without code changes. The API remains identical to the one used with CUDA.

python 
import torch

# Same code for CUDA (NVIDIA) and ROCm (AMD)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

x = torch.randn(1000, 1000, device=device)
y = torch.matmul(x, x.T)
print(f'Device: {y.device}')  # cuda:0 (NVIDIA or AMD)

Sources