Skip to content

Commit

Permalink
make nGPUs an optional parameter that if not specified runs with the …
Browse files Browse the repository at this point in the history
…max number of available GPUs
  • Loading branch information
Sophia Maedler committed Oct 26, 2023
1 parent 963ed89 commit 9cffaa1
Showing 1 changed file with 7 additions and 1 deletion.
8 changes: 7 additions & 1 deletion src/sparcscore/pipeline/segmentation.py
Original file line number Diff line number Diff line change
@@ -728,6 +728,9 @@ def process(self, input_image):
gc.collect() #perform garbage collection

#check that that number of GPUS is actually available
if "nGPUS" not in self.config.keys():
self.config["nGPUs"] = torch.cuda.device_count()

nGPUS = self.config["nGPUs"]
available_GPUs = torch.cuda.device_count()
processes_per_GPU = self.config["threads"]
@@ -1118,11 +1121,14 @@ def process(self):
self.log(f"A total of {len(segmentation_list)} processes need to be executed.")

#check that that number of GPUS is actually available
if "nGPUS" not in self.config.keys():
self.config["nGPUs"] = torch.cuda.device_count()

nGPUS = self.config["nGPUs"]
available_GPUs = torch.cuda.device_count()
processes_per_GPU = self.config["threads"]

if available_GPUs != self.config["nGPUs"]:
if available_GPUs < self.config["nGPUs"]:
self.log(f"Found {available_GPUs} but {nGPUS} specified in config.")

if available_GPUs >= 1:

0 comments on commit 9cffaa1

Please sign in to comment.