From 571bc20a1f32afded3711547e6c022ba2ca2c2bf Mon Sep 17 00:00:00 2001 From: MyGodItsFull0fStars Date: Wed, 6 Mar 2024 20:56:00 +0100 Subject: [PATCH] Added missing return type hint and updated doc string --- codecarbon/emissions_tracker.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/codecarbon/emissions_tracker.py b/codecarbon/emissions_tracker.py index 0e0beaf79..40eccc00c 100644 --- a/codecarbon/emissions_tracker.py +++ b/codecarbon/emissions_tracker.py @@ -485,11 +485,11 @@ def start_task(self, task_name=None) -> None: ) self._active_task = task_name - def stop_task(self, task_name: str = None) -> float: + def stop_task(self, task_name: str = None) -> EmissionsData: """ Stop tracking a dedicated execution task. Delta energy is computed by task, to isolate its contribution to total emissions. - :return: None + :return: EmissionData for an execution task """ task_name = task_name if task_name else self._active_task self._measure_power_and_energy() @@ -505,6 +505,7 @@ def stop_task(self, task_name: str = None) -> float: self._tasks[task_name].emissions_data = task_emission_data self._tasks[task_name].is_active = False self._active_task = None + return task_emission_data @suppress(Exception)