Why is device in device_dtype_mixin of type Union[str, torch.device]?

[daniellepintz]

Could anyone help me understand in which circumstances device would be a string?

https://github.com/PyTorchLightning/pytorch-lightning/blob/6e124e7207f6459cb43f540cfb5a1c6cc9b00f7a/pytorch_lightning/core/mixins/device_dtype_mixin.py#L39

cc @Borda who last committed this line

[Borda]

circumstances device would be a string?

you can have "cpu" or "cuda"

[daniellepintz]

thanks! and how/where would that be set? in the code above it looks like it will always be of type torch.device

[daniellepintz]

I am wondering in the context of XLAStatsMonitor:
https://github.com/PyTorchLightning/pytorch-lightning/blob/4bebf82f271fb577ed41831403551f0bff97651b/pytorch_lightning/callbacks/xla_stats_monitor.py#L78

where we pass pl_module.device to xm.get_memory_info which takes a string. This works because pl_module.device returns Union[str, torch.device], but I am wondering what is the string exactly? Can someone give me an example of what this string might equal here?

I am unable to test locally since I don't have access to TPUs.