From 21e51f6313ccd520e1f0af78e7f74ec649e8b80a Mon Sep 17 00:00:00 2001 From: AnnihilatorChess <150355206+AnnihilatorChess@users.noreply.github.com> Date: Mon, 17 Nov 2025 11:00:37 +0100 Subject: [PATCH 1/2] actually use the FNO blocks --- the_well/benchmark/models/fno/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/the_well/benchmark/models/fno/__init__.py b/the_well/benchmark/models/fno/__init__.py index 1773ac2a..7437f206 100755 --- a/the_well/benchmark/models/fno/__init__.py +++ b/the_well/benchmark/models/fno/__init__.py @@ -47,7 +47,7 @@ def forward(self, x: torch.Tensor, output_shape=None, **kwargs): x = self.domain_padding.pad(x) for layer_idx in range(self.n_layers): - self.optional_checkpointing( + x = self.optional_checkpointing( self.fno_blocks, x, layer_idx, output_shape=output_shape[layer_idx] ) From e4301602e57f2dddeaa87856dda49bbba47c663b Mon Sep 17 00:00:00 2001 From: AnnihilatorChess <150355206+AnnihilatorChess@users.noreply.github.com> Date: Mon, 17 Nov 2025 11:02:43 +0100 Subject: [PATCH 2/2] actually use the TFNO blocks --- the_well/benchmark/models/tfno/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/the_well/benchmark/models/tfno/__init__.py b/the_well/benchmark/models/tfno/__init__.py index 21d08518..31eb36a8 100755 --- a/the_well/benchmark/models/tfno/__init__.py +++ b/the_well/benchmark/models/tfno/__init__.py @@ -45,7 +45,7 @@ def forward(self, x: torch.Tensor, output_shape=None, **kwargs): x = self.domain_padding.pad(x) for layer_idx in range(self.n_layers): - self.optional_checkpointing( + x = self.optional_checkpointing( self.fno_blocks, x, layer_idx, output_shape=output_shape[layer_idx] )