search

funkelab / gunpowder

A library to facilitate machine learning on multi-dimensional images.
https://funkelab.github.io/gunpowder/
MIT License
79 stars 56 forks source link

Passing cuda device as an argument to Train, Predict nodes #188

Closed Mohinta2892 closed 10 months ago

Mohinta2892 commented 1 year ago

It would be helpful if we could pass a cuda device explicitly to the train and predict nodes for DGX systems, to be able to spawn multiple models on different cards simultaneously on a single machine. Currently both scripts default to "cuda:0" on device availability: "cuda:0"

I have added this into a fork of the repo, wanted to discuss if this is something useful to integrate here.

Best, Samia

pattonw commented 1 year ago

Yes, this would be quite helpful to integrate into gunpowder directly

Mohinta2892 commented 1 year ago

Thanks for replying Will! I shall raise a pull request soon-ish.

pattonw commented 10 months ago

solved in a pull request