Fix ZeusMonitor integration example

ml-energy · Nov 10, 2023 · bcfa4d2 · bcfa4d2
1 parent 4f207de
commit bcfa4d2
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/zeus/monitor/energy.py b/zeus/monitor/energy.py
@@ -89,9 +89,9 @@ class ZeusMonitor:
     result = monitor.end_window("entire_training")
 
     # Print the measurement result.
-    time_consumed, energy_consumed = prof_result.time, prof_result.energy
-    print(f"Training took {time_consumed} seconds.")
-    for gpu_idx, gpu_energy in zip(gpu_indices, energy_consumed):
+    print(f"Training took {result.time} seconds.")
+    print(f"Training consumed {result.total_energy} Joules.")
+    for gpu_idx, gpu_energy in result.energy.items():
         print(f"GPU {gpu_idx} consumed {gpu_energy} Joules.")
     ```