model cannot be trained on gpu

jiyuuchc / lacss

A deep learning model for single cell segmentation from microsopy images.

https://jiyuuchc.github.io/lacss/

MIT License

27 stars 4 forks source link

model cannot be trained on gpu #9

Closed xzx0554 closed 4 months ago

xzx0554 commented 4 months ago

Dear author, I'm using your lacss library version 0.8.1. When training a model on an A40 (48G) GPU, I noticed that the model only runs on the GPU during validation, but the training is done on the CPU. How can I resolve this issue? I would greatly appreciate your assistance.