Skip to content

set_device

torch.tvarant.set_device(device) → None

Sets the current device index.

Set the current device index.

Parameters:

Name Type Description Default
device int

Device ordinal. Must be in [0, device_count()).

 required

Examples:

>>> torch.tvarant.set_device(0)
Source code in torch_tvarant/tvarant/__init__.py 
85
86
87
88
89
90
91
92
93
94
def set_device(device: int) -> None:
    """Set the current device index.

    Args:
        device: Device ordinal. Must be in ``[0, device_count())``.

    Examples:
        >>> torch.tvarant.set_device(0)
    """
    _C.set_device(int(device))