mmengine.dist.get_data_device¶
- mmengine.dist.get_data_device(data)[源代码]¶
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.- 参数:
data (Tensor or Sequence or dict) – Inputs to be inferred the device.
- 返回:
The device of
data
.- 返回类型:
示例
>>> 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')