Device changed to CPU when calling torchinfo.summary() when model is created on MPS for Mac Users

[frankcaoyun]

Describe the bug
The model device will be changed to CPU when calling torchinfo.summary() when model is already created on MPS for macOS.

To Reproduce
Steps to reproduce the behavior:

1. On macOS, create a pytorch model and change the device to mps
2. Check the model device
3. Call torchinfo.summary() without passing in the device parameter
4. Check the model device again

Expected behavior
The code should handle this case automatically by detecting if the model is on mps for macOS

Desktop (please complete the following information):

• OS: macOS

Additional context
Appreciate your help.

[Tyler Yep]

HI, unfortunately I don't think we can always safely default to mps if it is available. If you run the tests in this project using mps as the device for all models and inputs, some tests fail:

        output_size = _unpool_output_size(input, kernel_size, _stride, padding, output_size)
>       return torch._C._nn.max_unpool2d(input, indices, output_size)
E       NotImplementedError: The operator 'aten::max_unpool2d' is not currently implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.

With PYTORCH_ENABLE_MPS_FALLBACK=1, we get errors still:

self = Linear(in_features=3, out_features=1, bias=True), input = tensor([1., 1., 1.])

    def forward(self, input: Tensor) -> Tensor:
>       return F.linear(input, self.weight, self.bias)
E       RuntimeError: Placeholder storage has not been allocated on MPS device!

Happy to reconsider this once feature parity is reached for mps, but this is a pretty big backwards-incompatible change as is. You can always specify the device if you want to use it.

[frankcaoyun]

Hi @TylerYep

Appreciate your update and understand the concern.

To not change the model device accidentally, will it be good to make a copy of the model to CPU by default? Other than the extra memory needed, will you foresee any other side effects? I think from a user's perspective, running a model summary shouldn't change the model device.

What do you think?

[Tyler Yep]

Yes, I think the correct behavior should be to not change the device of the model if device parameter is not specified, which is tracked in a separate issue #227

[frankcaoyun]

Got it. Look forward to the fix.