And then I guess this is the full function:
def create_freq_tensor(size, cf=.5**.5, device='cpu'):
    step=(cf/(size/2))
    end = (step * size)/2
    return torch.cat([torch.arange(0, end=end, step=step, device=device), torch.arange(-cf, end=0, step=step, device=device)])

def get_fft_scale(h, w, decay_power=.75, device='cpu'):
    d=.5**.5
    fy = create_freq_tensor(h, device=device)[:,None]
    if w % 2 == 1:
        wadd = 2
    else:
        wadd = 1
    fx = create_freq_tensor(w, device=device)[: w // 2 + wadd]        
    freqs = (fx*fx + fy*fy) ** decay_power
    scale = 1 / torch.max(freqs, torch.full_like(freqs, 1.0 / (max(w, h)*d)))
    return scale[None,None,...,None]

Is there a way I could do it better?