After calling summary(), model on CPU is pushed to GPU
ZhilinGuo opened this issue · comments
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:
- Create a Pytorch model instance on CPU, double check that
next(net.parameters()).is_cuda
returns False. - Run
summary()
on the model. - Check
next(net.parameters()).is_cuda
again. It now returns True. - 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.