Giters

TylerYep / torchinfo

View model summaries in PyTorch!

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

After calling summary(), model on CPU is pushed to GPU

ZhilinGuo opened this issue · comments

commented

Describe the bug
After calling summary(), the Pytorch model that was originally on CPU will be pushed to GPU.

To Reproduce
Steps to reproduce the behavior:

  1. Create a Pytorch model instance on CPU, double check that next(net.parameters()).is_cuda returns False.
  2. Run summary() on the model.
  3. Check next(net.parameters()).is_cuda again. It now returns True.
  4. Any operations with the model will give error unless model is manually transferred from GPU to CPU using to(device) (because other tensors are still on CPU).

Expected behavior
The model should stay on its original device.