Add check for input dimension

Project.toml

@@ -8,19 +8,21 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 
 [compat]
-Tullio = "0.3"
-ChainRules = "1"
 CUDA = "5"
-Lux = "1"
+ChainRules = "1"
 ComponentArrays = "0.15"
-LuxCore = "1"
 FFTW = "1"
+JuliaFormatter = "1.0.62"
+Lux = "1"
+LuxCore = "1"
 NNlib = "0.9"
+Tullio = "0.3"
 julia = "1.10"

src/models.jl

@@ -280,6 +280,12 @@ function ((;)::CNO)(x, params, state)
     bottleneck_ranges = state.bottleneck_ranges
     reversed_bottleneck_ranges = state.reversed_bottleneck_ranges
 
+    # Assert that the dimensions of x are correct
+    for i in range(1, D)
+        @assert size(x, i) == N "ERROR: x has dimension $i = $(size(x, i)) but it should be $N"
+    end
+
+
     # we have to keep track of each downsampled state
     intermediate_states = []