Source code for utils.device
[docs]def get_device(tensor):
"""Get the device of a tensor
Parameters
----------
tensor: Tensor
tensor from which to extract the device
Returns
-------
computation device (0-n or cpu)
"""
device = tensor.get_device()
return "cpu" if 0 > device else device
[docs]def to_same_device(*tensors, device=None):
"""Move all tensors to the same device (by default the first tensor's one)
Parameters
----------
tensors: Tensor
tensors which we want to move to a common device
Returns
-------
a tuple of tensors in the same order as the arguments moved to a common device
"""
if device is None:
device = get_device(tensors[0])
return tuple(tensor.to(device) for tensor in tensors)