fix type hints

Signed-off-by: nstarman <nstarman@users.noreply.github.com>
GalOrrery · Jul 17, 2023 · 1bc9918 · 1bc9918
1 parent 176a468
commit 1bc9918
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/stream_ml/pytorch/_multi.py b/src/stream_ml/pytorch/_multi.py
@@ -106,7 +106,7 @@ class MixtureModel(ModelsBase, CoreMixtureModel[Array, NNModel]):
         The neural network that is used to combine the components.
     """
 
-    net: NNField[NNModel] = NNField(default=MISSING)
+    net: NNField[NNModel, NNModel] = NNField(default=MISSING)
 
     def __new__(cls: type[Self], *args: Any, **kwargs: Any) -> Self:  # noqa: ARG003
         """Initialize the model. This is needed for PyTorch."""

diff --git a/src/stream_ml/pytorch/builtin/_isochrone.py b/src/stream_ml/pytorch/builtin/_isochrone.py
@@ -23,10 +23,10 @@
 if TYPE_CHECKING:
     from scipy.interpolate import CubicSpline
 
-    from stream_ml.core.typing import BoundsT, NNModel
+    from stream_ml.core.typing import BoundsT
 
     from stream_ml.pytorch.params import Params
-    from stream_ml.pytorch.typing import Array, ArrayNamespace
+    from stream_ml.pytorch.typing import Array, ArrayNamespace, NNModel
 
 dm_sigma_const: Final = 5 / xp.log(xp.asarray(10))
 # Constant for first order error propagation of parallax -> distance modulus