chore: mypy typing fix

instadeepai · Feb 27, 2024 · c973032 · c973032
1 parent d2537cf
commit c973032
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/marl_eval/plotting_tools/plot_utils.py b/marl_eval/plotting_tools/plot_utils.py
@@ -19,6 +19,7 @@
 import numpy as np
 import seaborn as sns
 from matplotlib.axes import Axes
+from matplotlib.figure import Figure
 from rliable.plot_utils import _annotate_and_decorate_axis
 
 
@@ -36,7 +37,7 @@ def plot_single_task_curve(
     legend_map: Optional[Dict] = None,
     run_times: Optional[Dict] = None,
     **kwargs: Any,
-) -> Axes:
+) -> Figure:
     """Plots an aggregate metric with CIs as a function of environment frames.
 
     Args: