Update worker_training.py

C-Achard · C-Achard · commit 09cc3401a079 · 2023-07-19T16:58:55.000+02:00
diff --git a/napari_cellseg3d/code_models/worker_training.py b/napari_cellseg3d/code_models/worker_training.py
@@ -546,10 +546,10 @@ def get_loader_func(num_samples):
                 if device.type == "cuda":
                     self.log("Memory Usage:")
                     alloc_mem = round(
-                        torch.cuda.memory_allocated(device_id) / 1024**3, 1
+                        torch.cuda.memory_allocated(device) / 1024**3, 1
                     )
                     reserved_mem = round(
-                        torch.cuda.memory_reserved(device_id) / 1024**3, 1
+                        torch.cuda.memory_reserved(device) / 1024**3, 1
                     )
                     self.log(f"Allocated: {alloc_mem}GB")
                     self.log(f"Cached: {reserved_mem}GB")

Original file line number	Diff line number	Diff line change
`@@ -546,10 +546,10 @@ def get_loader_func(num_samples):`
`546`	`546`	`if device.type == "cuda":`
`547`	`547`	`self.log("Memory Usage:")`
`548`	`548`	`alloc_mem = round(`
`549`		`- torch.cuda.memory_allocated(device_id) / 1024**3, 1`
	`549`	`+ torch.cuda.memory_allocated(device) / 1024**3, 1`
`550`	`550`	`)`
`551`	`551`	`reserved_mem = round(`
`552`		`- torch.cuda.memory_reserved(device_id) / 1024**3, 1`
	`552`	`+ torch.cuda.memory_reserved(device) / 1024**3, 1`
`553`	`553`	`)`
`554`	`554`	`self.log(f"Allocated: {alloc_mem}GB")`
`555`	`555`	`self.log(f"Cached: {reserved_mem}GB")`