Models on cuda are always moved to GPU 0 by torchinfo.summary

[Zexi]

Describe the bug
For CPU usage, torchinfo.summary works well.
For models on GPU, if the model were not on GPU 0 (e.g. on GPU 5), torchinfo.summary will move it to GPU 0.

To Reproduce
Steps to reproduce the behavior:

• execute the python script in command line to prevent python from terminating

from torch import nn
import torch
from torchinfo import summary

linear = nn.Linear(1000, 1000).to(1)
a = torch.randn(32, 1000).to(1)
summary(linear, input_data=a)

• I encountered 2 cases:

1. Run nvidia-smi, you will see on GPU 0 a python process
2. Replacing nn.Linear with torchaudio.transforms.MelSpectrogram or something similar, it will report that the model and the data are on different devices, which prevents the computation.

Expected behavior
summary should be executed on the given card

Desktop (please complete the following information):

• OS: macOS 13.1, Ubuntu 16.04
• torchinfo 1.7.1

[Tyler Yep]

Have you tried using the device parameter to summary?

[Sebastian Müller]

I have sucessfully replicated the error.

Running

from torch import nn
import torch
from torchinfo import summary

 
linear = nn.Linear(1000, 1000).to(1)
a = torch.randn(32, 1000).to(1)
_ = summary(linear, input_data=a)
next(linear.parameters()).device

yields:

device(type='cuda', index=0)

And it turns out that using the device parameter works just fine:

from torch import nn
import torch
from torchinfo import summary

linear = nn.Linear(1000, 1000).to(1)
a = torch.randn(32, 1000).to(1)
_ = summary(linear, input_data=a, device=torch.device("cuda:1"))
next(linear.parameters()).device

yields:

device(type='cuda', index=1)

Still, this issue is adressed in pull-request#211.