Pass backend name to EnergyTracker in Training scenario (huggingface#279

)
vicoooo26 · Oct 23, 2024 · c3908cb · c3908cb
1 parent 561deca
commit c3908cb
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/optimum_benchmark/scenarios/training/scenario.py b/optimum_benchmark/scenarios/training/scenario.py
@@ -54,7 +54,9 @@ def run(self, backend: Backend[BackendConfigT]) -> BenchmarkReport:
 
         if self.config.energy:
             self.logger.info("\t+ Creating energy tracking context manager")
-            energy_tracker = EnergyTracker(device=backend.config.device, device_ids=backend.config.device_ids)
+            energy_tracker = EnergyTracker(
+                device=backend.config.device, backend=backend.config.name, device_ids=backend.config.device_ids
+            )
 
         if self.config.memory:
             self.logger.info("\t+ Entering memory tracking context manager")

diff --git a/tests/test_api.py b/tests/test_api.py
@@ -64,7 +64,13 @@ def test_api_launch(device, scenario, library, task, model):
 
     if scenario == "training":
         if library == "transformers":
-            scenario_config = TrainingConfig(memory=True, latency=True, warmup_steps=2, max_steps=5)
+            scenario_config = TrainingConfig(
+                memory=True,
+                latency=True,
+                energy=not is_rocm_system(),
+                warmup_steps=2,
+                max_steps=5,
+            )
         else:
             pytest.skip("Training scenario is only available for Transformers library")