search

TylerYep / torchinfo

View model summaries in PyTorch!
MIT License
2.56k stars 119 forks source link

CPU vs GPU usage #202

Closed abs-xyz closed 1 year ago

abs-xyz commented 1 year ago

Is your feature request related to a problem? Please describe. I was facing the issue described here. I wanted to train on CPU, although CUDA was available. It took me a long time to figure out that the problem was due to my incorrect usage of summary.

Describe the solution you'd like It would be great if the forward pass won't lead to the same model weights being put on the GPU. Maybe the model instance can be cloned and then this forward pass be done? This would not disrupt the further training, even if someone forgets to set the device argument.

TylerYep commented 1 year ago

Perhaps torchinfo can catch the exception and suggest a fix using device. I'm not very interested in cloning the model instance to support this mistake, since it seems like it is easily avoided by reading the documentation.