Default to same device as provided model

sksq96 / pytorch-summary

Model summary in PyTorch similar to `model.summary()` in Keras

MIT License

4.01k stars 412 forks source link

Default to same device as provided model #135

Closed m3at closed 2 months ago

m3at commented 4 years ago

Hello, thanks for working on pytorch-summary. Models are often created on CPU, or not on the first cuda device. Defaulting to the current model device might be helpful.

nijfranck commented 3 years ago

Defaulting to current model device and number is definitely more flexible. Thanks for the pull request and commit @m3at