FNO recompiles for every new input resolution with torch.compile(dynamic=True) #2207

Compiling FNO with dynamic=True still recompiles every time the grid size changes. With a 2D FNO and inputs at 16/24/32/40, I get 4 graphs instead of 1, and the recompile reason is tensor 'x' size mismatch at index 2. TORCH_LOGS=dynamic shows the spatial size gets specialized at torch.linspace(0, 1, size_x, ...) in the grid encoders' _meshgrid in nn/module/fno_layers.py (1D/2D/3D/4D). Since FNO is meant to run at different resolutions, each new size costs a recompile, and the default compile falls back to eager after the recompile limit.

Environment: torch 2.13.0+cu130, RTX 4090, main ebf46ee.

Building the grid with torch.arange(n) / max(n - 1, 1) keeps the size dynamic, and 1D/2D then compile a single graph that matches eager. 3D then hits an Inductor ValueRangeError inside PyTorch, so it may need to stay as-is for now. Happy to open a PR.

Repro:

# FNO recompiles for every new grid size under torch.compile(dynamic=True)
import torch
from physicsnemo.models.fno import FNO
from torch._dynamo.utils import counters
m = FNO(in_channels=3, out_channels=2, dimension=2, latent_channels=16, num_fno_modes=8).cuda().eval()
cm = torch.compile(m, dynamic=True, fullgraph=True)
with torch.no_grad():
    for n in (16, 24, 32, 40):
        cm(torch.randn(2, 3, n, n, device="cuda"))
print("unique graphs:", counters["stats"]["unique_graphs"])  # 4 on main, expected 1