mypy

invrs-io · Nov 17, 2023 · 180b97a · 180b97a
1 parent 7fb8020
commit 180b97a
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/src/invrs_gym/challenges/sorter/common.py b/src/invrs_gym/challenges/sorter/common.py
@@ -230,7 +230,9 @@ def response(
         spec = dataclasses.replace(
             self.spec,
             thickness_cap=params[THICKNESS_CAP],  # type: ignore[arg-type]
-            thickness_metasurface=params[THICKNESS_METASURFACE],  # type: ignore[arg-type]
+            thickness_metasurface=(
+                params[THICKNESS_METASURFACE]  # type: ignore[arg-type]
+            ),
             thickness_spacer=params[THICKNESS_SPACER],  # type: ignore[arg-type]
         )
         return simulate_sorter(
@@ -375,9 +377,9 @@ def simulate_sorter(
 
     layer_thicknesses = [
         jnp.zeros(()),  # Ambient
-        spec.thickness_cap.array,
-        spec.thickness_metasurface.array,
-        spec.thickness_spacer.array,
+        jnp.asarray(spec.thickness_cap.array),
+        jnp.asarray(spec.thickness_metasurface.array),
+        jnp.asarray(spec.thickness_spacer.array),
         jnp.asarray(spec.offset_monitor_substrate),  # Substrate
     ]