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 in data 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.

返回类型:

torch.device

示例

>>> 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')