chore: Fix for tests

neptune-ai · Aug 19, 2024 · 1902e67 · 1902e67
1 parent f62d559
commit 1902e67
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/src/neptune/common/hardware/gpu/gpu_monitor.py b/src/neptune/common/hardware/gpu/gpu_monitor.py
@@ -46,7 +46,7 @@ def get_card_used_memory_in_bytes(self, card_index):
         return self.__nvml_get_or_else(lambda: nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(card_index)).used)
 
     def get_card_power_usage(self, card_index):
-        return self.__nvml_get_or_else(lambda: nvmlDeviceGetPowerUsage(nvmlDeviceGetHandleByIndex(card_index)).used)
+        return self.__nvml_get_or_else(lambda: nvmlDeviceGetPowerUsage(nvmlDeviceGetHandleByIndex(card_index)))
 
     def get_card_max_power_rating(self):
         def read_max_power_rating():

diff --git a/tests/unit/neptune/legacy/internal/hardware/gauges/test_gpu_gauges.py b/tests/unit/neptune/legacy/internal/hardware/gauges/test_gpu_gauges.py
@@ -83,9 +83,7 @@ def test_gpu_power_gauge(self, nvmlDeviceGetPowerUsage):
         # and
         gpu_power_info = MagicMock()
         gpu_power_info.used = 15.0 * MILLIWATTS_IN_ONE_WATT
-        nvmlDeviceGetPowerUsage.side_effect = lambda handle: (
-            gpu_power_info if handle == self.gpu_card_handle else None
-        )
+        nvmlDeviceGetPowerUsage.return_value = gpu_power_info.used
 
         # when
         power_watts = gauge.value()