Merge pull request neuraloperator#177 from neuraloperator/dev/colin

fix SFNO example
ziqi-ma · Jul 12, 2023 · 5efc112 · 5efc112
2 parents 236d582 + e1dcb45
commit 5efc112
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/examples/plot_SFNO_swe.py b/examples/plot_SFNO_swe.py
@@ -113,7 +113,8 @@
     # Ground-truth
     y = data['y'][0, ...].numpy()
     # Model prediction
-    out = model(x.unsqueeze(0)).squeeze()[0, ...].detach().numpy()
+    x_in = x.unsqueeze(0).to(device)
+    out = model(x_in).squeeze()[0, ...].detach().cpu().numpy()
     x = x[0, ...].detach().numpy()
 
     ax = fig.add_subplot(2, 3, index*3 + 1)

diff --git a/neuralop/models/spherical_convolution.py b/neuralop/models/spherical_convolution.py
@@ -164,7 +164,7 @@ def get_contract_fun(weight, implementation='reconstructed', separable=False):
 
 class FactorizedSphericalConv(nn.Module):
     def __init__(self, in_channels, out_channels, n_modes, incremental_n_modes=None, bias=True,
-                 n_layers=1, separable=False, output_scaling_factor=None,
+                 n_layers=1, separable=False, output_scaling_factor=None, fno_block_precision='full',
                  rank=0.5, factorization='cp', implementation='reconstructed', 
                  fixed_rank_modes=False, joint_factorization=False, decomposition_kwargs=dict(),
                  init_std='auto', fft_norm='backward'):