src/ptbench/data/dataset.py

@@ -333,12 +333,12 @@ class CachedDataset(torch.utils.data.Dataset):
 
         logger.info(f"Caching {self.subset} samples")
         for sample in tqdm(self._samples):
-            sample["data"] = self.transforms(
-                self.raw_data_loader(sample["data"])
-            )
+            sample["data"] = self.raw_data_loader(sample["data"])
 
     def __getitem__(self, idx):
-        return self._samples[idx]
+        sample = self._samples[idx].copy()
+        sample["data"] = self.transforms(sample["data"])
+        return sample
 
     def __len__(self):
         return len(self._samples)

src/ptbench/utils/resources.py

@@ -100,7 +100,7 @@ def gpu_constants():
     """
     return run_nvidia_smi(
         ("gpu_name", "driver_version", "memory.total"),
-        ("gpu_name", "gpu_driver_version", "gpu_memory_total"),
+        ("gpu_name", "gpu_driver_version", "gpu_memory_total_GB"),
     )
 
 
@@ -135,10 +135,10 @@ def gpu_log():
             "utilization.gpu",
         ),
         (
-            "gpu_memory_total",
-            "gpu_memory_used",
-            "gpu_memory_free",
-            "gpu_percent",
+            "gpu_memory_total_GB",
+            "gpu_memory_used_GB",
+            "gpu_memory_free_percent",
+            "gpu_usage_percent",
         ),
     )
 
@@ -165,7 +165,7 @@ def cpu_constants():
         1. ``cpu_count`` (:py:class:`int`): number of logical CPUs available
     """
     return (
-        ("cpu_memory_total", psutil.virtual_memory().total / GB),
+        ("cpu_memory_total_GB", psutil.virtual_memory().total / GB),
         ("cpu_count", psutil.cpu_count(logical=True)),
     )
 
@@ -245,9 +245,9 @@ class CPULogger:
                 # at this point, but ensures to update counts later on
                 gone.add(k)
         return (
-            ("cpu_memory_used", psutil.virtual_memory().used / GB),
-            ("cpu_rss", sum([k.rss for k in memory_info]) / GB),
-            ("cpu_vms", sum([k.vms for k in memory_info]) / GB),
+            ("cpu_memory_used_GB", psutil.virtual_memory().used / GB),
+            ("cpu_rss_GB", sum([k.rss for k in memory_info]) / GB),
+            ("cpu_vms_GB", sum([k.vms for k in memory_info]) / GB),
             ("cpu_percent", sum(cpu_percent)),
             ("cpu_processes", len(self.cluster) - len(gone)),
             ("cpu_open_files", sum(open_files)),