import torch


def tensor_shift_fft(fft):
    """
    Shift zero frequency spatial frequency component to center of 2D image. For Pytorch implementation
    Args:
        fft: 2D FFT obtained using torch_fft function
    Returns:
        shifted FFT with DC frequency component in center of image.
    """
    m, n = fft.shape
    out = torch.cat((fft[-m//2:], fft[:-m//2]), dim=0)
    return torch.cat((out[:, -n//2:], out[:, :-n//2]), dim=1)