mmengine.dist.cast_data_device

mmengine.dist.cast_data_device(data, device, out=None)

Recursively convert Tensor in data to device.

If data has already on the device, it will not be casted again.

参数:

• data (Tensor or list or dict) – Inputs to be casted.

• device (torch.device) – Destination device type.

• out (Tensor or list or dict, optional) – If out is specified, its value will be equal to data. Defaults to None.

返回:

data was casted to device.

返回类型:

Tensor or list or dict