append einsum and 2dsolver version of EinFFT

kyegomez · Mar 25, 2024 · 44a4c21 · 44a4c21
1 parent b0b2652
commit 44a4c21
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/simba/main.py b/simba/main.py
@@ -83,11 +83,14 @@ def forward(self, x: Tensor) -> Tensor:
         iffted = torch.fft.ifft2(activated_complex, dim=(-2, -1))
         return iffted.real
 
-
-x = torch.randn(1, 3, 64, 64)  # Random input tensor
-einfft = EinFFT(3, 64, 64)  # Instantiate EinFFT module
-output = einfft(x)  # Apply EinFFT
-print(output)  # Print module architecture
+# Random input tensor
+x = torch.randn(1, 3, 64, 64)
+# Instantiate EinFFT module
+einfft = EinFFT(3, 64, 64)
+# Apply EinFFT to get an output
+output = einfft(x)
+# Print output tensor
+print(output)
 
 
 class Simba(nn.Module):