mmengine.dist.get_data_device¶
- mmengine.dist.get_data_device(data)[source]¶
Return the device of
data
.If
data
is a sequence of Tensor, all items indata
should have a same device type.If
data
is a dict whose values are Tensor, all values should have a same device type.- Parameters
data (Tensor or Sequence or dict) – Inputs to be inferred the device.
- Returns
The device of
data
.- Return type
Examples
>>> import torch >>> from mmengine.dist import cast_data_device >>> # data is a Tensor >>> data = torch.tensor([0, 1]) >>> get_data_device(data) device(type='cpu') >>> # data is a list of Tensor >>> data = [torch.tensor([0, 1]), torch.tensor([2, 3])] >>> get_data_device(data) device(type='cpu') >>> # data is a dict >>> data = {'key1': torch.tensor([0, 1]), 'key2': torch.tensor([0, 1])} >>> get_data_device(data) device(type='cpu')