From 180b97ae86811a7d2d3a4f5e623b30cee4018db5 Mon Sep 17 00:00:00 2001 From: Martin Schubert Date: Fri, 17 Nov 2023 14:43:33 -0800 Subject: [PATCH] mypy --- src/invrs_gym/challenges/sorter/common.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/invrs_gym/challenges/sorter/common.py b/src/invrs_gym/challenges/sorter/common.py index 256ee57..10d028c 100644 --- 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 ]