actually use the TFNO blocks

AnnihilatorChess · web-flow · commit e4301602e57f · 2025-11-17T11:02:43.000+01:00
diff --git 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]
             )
 

Original file line number	Diff line number	Diff line change
`@@ -45,7 +45,7 @@ def forward(self, x: torch.Tensor, output_shape=None, **kwargs):`
`45`	`45`	`x = self.domain_padding.pad(x)`
`46`	`46`
`47`	`47`	`for layer_idx in range(self.n_layers):`
`48`		`- self.optional_checkpointing(`
	`48`	`+ x = self.optional_checkpointing(`
`49`	`49`	`self.fno_blocks, x, layer_idx, output_shape=output_shape[layer_idx]`
`50`	`50`	`)`
`51`	`51`