search

DevoLearn / devolearn

Accelerate data driven research on embryos with deep learning models :microscope:
MIT License
43 stars 30 forks source link

fix #4 : add gpu support #34

Closed rtharungowda closed 3 years ago

rtharungowda commented 3 years ago

Add gpu/cuda support for models embryo_generator_model, embryo_segmentor and lineage_population_model

Mayukhdeb commented 3 years ago

@rtharungowda thanks for the PR, did you test it out on a cuda capable device ?

rtharungowda commented 3 years ago

@rtharungowda thanks for the PR, did you test it out on a cuda capable device ?

Yes, I tested it on colab https://colab.research.google.com/drive/1AcxKTV-iCAAMpaSqnLiKkIrTIb4Nh__l?usp=sharing

Mayukhdeb commented 3 years ago

Looks good. Can you make one more minor edit ?

let us stick to the device argument for selecting the device.

It should look like: 

segmentor = some_model(..., device = 'cuda')

or 

segmentor = some_model(..., device = 'cpu')

And then set the device as self.device = device