search

juglab / cryoCARE_pip

PIP package of cryoCARE
BSD 3-Clause "New" or "Revised" License
25 stars 14 forks source link

Added GPU ID selection capability #24

Closed EuanPyle closed 2 years ago

EuanPyle commented 2 years ago

Hi, I've added a few lines to allow the user to select GPU IDs in train.py and predict.py. The user specifies GPU IDs as a list or an integer in the json config files for those scripts. Importantly, I've also added a line to set_memory_growth to true for GPUs as without it I was consistently getting OOM errors if you think this is a sensible way to deal with that? The structure of this PR could probably do with some work too: I wasn't sure exactly where to place the function and how to call it within the current cryoCARE structure so I simply added a function in predict.py and also call it in train.py Thanks, Euan

thorstenwagner commented 2 years ago

Looks all fine for me. Sorry, I totally forgot to look at this. I will merge it.

EuanPyle commented 2 years ago

Thanks!