From b00aaaed185e2f00e84a34986034c54109a52084 Mon Sep 17 00:00:00 2001
From: Angela <angela@cleanlab.ai>
Date: Mon, 17 Jun 2024 09:58:57 -0700
Subject: [PATCH] mypy

---
 cleanlab_studio/studio/studio.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cleanlab_studio/studio/studio.py b/cleanlab_studio/studio/studio.py
index 1416e1b8..a7fa4884 100644
--- a/cleanlab_studio/studio/studio.py
+++ b/cleanlab_studio/studio/studio.py
@@ -379,7 +379,7 @@ def download_pred_probs(
         If you want to work with predicted probabilities for an image project, the recommended workflow is to download probabilities with the option `keep_id=True`, and then do a join with the original dataset on the ID column.
         Alternatively, you can follow the steps [here](/reference/python/studio#method-download_embeddings), and filter out the rows that were not analyzed. The filtered dataset will then have rows that align with the predicted probabilities DataFrame.
         """
-        pred_probs: Union[npt.NDArray[np.float_], pd.DataFrame] = api.download_array(
+        pred_probs: Union[npt.NDArray[np.float64], pd.DataFrame] = api.download_array(
             self._api_key, cleanset_id, "pred_probs"
         )
         if not isinstance(pred_probs, pd.DataFrame):
@@ -396,7 +396,7 @@ def download_pred_probs(
     def download_embeddings(
         self,
         cleanset_id: str,
-    ) -> npt.NDArray[np.float_]:
+    ) -> npt.NDArray[np.float64]:
         """
         Downloads feature embeddings for a cleanset (available only for text and image projects).
         These are numeric vectors produced via neural network representations of each data point in your dataset.