torchinfo.summary should not change the device of model

[gitxw]

Describe the bug
torchinfo.summary should not change the device of model

To Reproduce
model = model.cpu()
torchinfo.summary(model)
why model was changed to cuda after calling torchinfo.summary() if I have a GraphicCard, It was Unexpected. I just want to dump the structure of my model.

Expected behavior
torchinfo.summary should keep the status of model, not change it.
the param "device" of the function torchinfo.summary is not necessary.

Desktop (please complete the following information):

• OS: Windows11

[1-ashraful-islam]

Had same issue using torchinfo.summary Current workaround seems to be passing device as additional argument.

Code example:

use_cuda = torch.cuda.is_available()
use_mps = torch.backends.mps.is_available()

if use_cuda:
    device = torch.device("cuda")
elif use_mps:
    device = torch.device("mps")
else:
    device = torch.device("cpu")

print(f"Using device (CUDA/MPS/CPU): {device}")

# initialize the model
model = ConvModel().to(device)

# print model info
print(summary(model, input_size=(batch_size, 1, 28, 28), device = device))

Device: Mac with apple silicon

[EbaraKoji]

I encountered the same problem on my M1 Mac.

The culprit seems to be get_device() function called in the summary().
This function does not consider mps device and device is set as cpu even when model parameters' device is mps.