Merge pull request #64 from AnnihilatorChess/master

mikemccabe210 · web-flow · commit 25bda95f11b4 · 2025-11-17T10:28:40.000-05:00
FNO and TFNO did not the FNO/TFNO blocks
diff --git 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]
             )
 
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
`@@ -47,7 +47,7 @@ def forward(self, x: torch.Tensor, output_shape=None, **kwargs):`
`47`	`47`	`x = self.domain_padding.pad(x)`
`48`	`48`
`49`	`49`	`for layer_idx in range(self.n_layers):`
`50`		`- self.optional_checkpointing(`
	`50`	`+ x = self.optional_checkpointing(`
`51`	`51`	`self.fno_blocks, x, layer_idx, output_shape=output_shape[layer_idx]`
`52`	`52`	`)`
`53`	`53`
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`