Fixed dimension of frequency threshold in 2d

[scaomath]

Based on the paper's Figure 3: https://openreview.net/pdf?id=EXHG-A3jlM

I think the number of modes threshold in 2d

AFNO-transformer/afno/afno2d.py

Lines 54 to 55 in 941e924

	total_modes = N // 2 + 1
	kept_modes = int(total_modes * self.hard_thresholding_fraction)

should be W instead of N since N = H*W, similar applies to bfno2d.py.

Nice paper btw. @jtguibas