Fix: Add getAverageMemoryPowerUsage to GPUs as well (#124)

ml-energy · Sep 10, 2024 · 3693d5a · 3693d5a
1 parent fcd6abc
commit 3693d5a
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/zeus/device/gpu/common.py b/zeus/device/gpu/common.py
@@ -253,6 +253,10 @@ def getInstantPowerUsage(self, gpu_index: int) -> int:
         """Return the current power draw of the GPU. Units: mW."""
         return self.gpus[gpu_index].getInstantPowerUsage()
 
+    def getAverageMemoryPowerUsage(self, gpu_index: int) -> int:
+        """Return the average power usage of the GPU's memory. Units: mW."""
+        return self.gpus[gpu_index].getAverageMemoryPowerUsage()
+
     def supportsGetTotalEnergyConsumption(self, gpu_index: int) -> bool:
         """Check if the GPU supports retrieving total energy consumption."""
         return self.gpus[gpu_index].supportsGetTotalEnergyConsumption()