Fix typing bug

.pre-commit-config.yaml

@@ -22,7 +22,8 @@ repos:
     rev: v0.991
     hooks:
       - id: mypy
-        args: [
+        args:
+          [
             "--ignore-missing-imports",
             "--scripts-are-modules",
             "--disallow-incomplete-defs",
@@ -35,5 +36,5 @@ repos:
             "--disallow-untyped-calls",
             "--install-types",
             "--non-interactive",
-            "--follow-imports=skip", # This is temporary until the mbi directory is not excluded
+            "--follow-imports=skip",
           ]

src/fdiff/utils/fourier.py

@@ -78,4 +78,9 @@ def idft(x: torch.Tensor) -> torch.Tensor:
     # Apply IFFT
     x_time = irfft(x_freq, dim=1, norm="ortho")
 
+    assert isinstance(x_time, torch.Tensor)
+    assert (
+        x_time.size() == x.size()
+    ), f"The inverse DFT and the input should have the same size. Got {x_time.size()} and {x.size()} instead."
+
     return x_time