fix: weird type error

Author: hack-hard

src/models.jl

@@ -3,10 +3,15 @@ using Lux
 function select_and_preprocess((point, atoms); cutoff_radius)
     select_and_preprocess(point, atoms; cutoff_radius)
 end
-function select_and_preprocess(point::Batch, atoms::AnnotedKDTree{Sphere{T}}; cutoff_radius) where T
+function select_and_preprocess(
+        point::Batch, atoms::AnnotedKDTree{Sphere{T}}; cutoff_radius) where {T}
     neighboord = Folds.map(point.field) do point
-		select_neighboord(point, atoms; cutoff_radius)::StructVector{Sphere{T}}
-	end |> Batch{Vector{<:StructVector{Sphere{T}}}}
+        select_neighboord(point,
+            atoms;
+            cutoff_radius)::StructVector{
+            Sphere{T}, @NamedTuple{center::Vector{Point3{T}}, r::Vector{T}}, Int64}
+    end |> Batch{Vector{StructVector{
+            Sphere{T}, @NamedTuple{center::Vector{Point3{T}}, r::Vector{T}}, Int64}}}
     preprocessing((point, neighboord))
 end
 
@@ -19,9 +24,9 @@ function evaluate_if_atoms_in_neighboord(layer, arg::AbstractArray, ps, st; zero
 end
 
 function general_angular_dense(main_chain, secondary_chain; name::String,
-    van_der_waals_channel=false, on_gpu=true, cutoff_radius::Float32=3.0f0)
+        van_der_waals_channel = false, on_gpu = true, cutoff_radius::Float32 = 3.0f0)
     main_chain = DeepSet(Chain(
-        symetrise(; cutoff_radius, device=on_gpu ? gpu_device() : identity),
+        symetrise(; cutoff_radius, device = on_gpu ? gpu_device() : identity),
         main_chain
     ))
     function add_van_der_waals_channel(main_chain)
@@ -41,7 +46,7 @@ end
 	`tiny_angular_dense` is a function that generate a lux model.
 
 """
-function tiny_angular_dense(; van_der_waals_channel=false, kargs...)
+function tiny_angular_dense(; van_der_waals_channel = false, kargs...)
     general_angular_dense(
         Parallel(.*,
             Chain(Dense(6 => 7, elu),
@@ -52,12 +57,12 @@ function tiny_angular_dense(; van_der_waals_channel=false, kargs...)
             BatchNorm(4 + van_der_waals_channel),
             Dense(4 + van_der_waals_channel => 6, elu),
             Dense(6 => 1, sigmoid_fast));
-        name="tiny_angular_dense_" *
-             (van_der_waals_channel ? "v" : ""),
+        name = "tiny_angular_dense_" *
+               (van_der_waals_channel ? "v" : ""),
         van_der_waals_channel, kargs...)
 end
 
-function light_angular_dense(; van_der_waals_channel=false, kargs...)
+function light_angular_dense(; van_der_waals_channel = false, kargs...)
     general_angular_dense(
         Parallel(.*,
             Chain(Dense(6 => 10, elu),
@@ -68,13 +73,13 @@ function light_angular_dense(; van_der_waals_channel=false, kargs...)
             BatchNorm(5 + van_der_waals_channel),
             Dense(5 + van_der_waals_channel => 10, elu),
             Dense(10 => 1, sigmoid_fast));
-        name="light_angular_dense_" *
-             (van_der_waals_channel ? "v" : ""),
+        name = "light_angular_dense_" *
+               (van_der_waals_channel ? "v" : ""),
         van_der_waals_channel, kargs...)
 end
 
 function medium_angular_dense(;
-    van_der_waals_channel=false, kargs...)
+        van_der_waals_channel = false, kargs...)
     general_angular_dense(
         Parallel(.*,
             Chain(Dense(6 => 15, elu),
@@ -83,11 +88,11 @@ function medium_angular_dense(;
         ),
         Chain(
             BatchNorm(10 + van_der_waals_channel),
-            Dense(10 + van_der_waals_channel => 5; use_bias=false),
+            Dense(10 + van_der_waals_channel => 5; use_bias = false),
             Dense(5 => 10, elu),
             Dense(10 => 1, sigmoid_fast));
-        name="medium_angular_dense_" *
-             (van_der_waals_channel ? "v" : ""),
+        name = "medium_angular_dense_" *
+               (van_der_waals_channel ? "v" : ""),
         van_der_waals_channel,
         kargs...)
 end
@@ -113,3 +118,7 @@ function get_cutoff_radius(x::Lux.AbstractExplicitLayer)
     get_preprocessing(x).fun.kargs[:cutoff_radius]
 end
 get_cutoff_radius(x::Lux.StatefulLuxLayer) = get_cutoff_radius(x.model)
+function get_cutoff_radius(x::Lux.AbstractExplicitLayer)
+    get_preprocessing(x).fun.kargs[:cutoff_radius]
+end
+get_cutoff_radius(x::Lux.StatefulLuxLayer) = get_cutoff_radius(x.model)