fix types

pymc-devs · Sep 20, 2023 · 2bf15ff · 2bf15ff
1 parent 84b29f7
commit 2bf15ff
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pymc_bart/utils.py b/pymc_bart/utils.py
@@ -536,7 +536,7 @@ def _get_axes(
             fig = ax[0].get_figure()
         else:
             axes = [ax]
-            fig = ax.get_figure()
+            fig = ax.get_figure()  # type: ignore
 
     return fig, axes, shape
 
@@ -718,7 +718,7 @@ def plot_variable_importance(
     figsize: Optional[Tuple[float, float]] = None,
     samples: int = 100,
     random_seed: Optional[int] = None,
-) -> Tuple[npt.NDArray[np.int_], List[plt.axes]]:
+) -> Tuple[npt.NDArray[np.int_], List[plt.Axes]]:
     """
     Estimates variable importance from the BART-posterior.