How to use 'torch check gpu' in Python

95
@staticmethod
96
def check_device(cuda: Union[bool, str, torch.device]) -> torch.device:
97
    if isinstance(cuda, bool):
98
        if cuda:
99
            if torch.cuda.is_available():
100
                return torch.device('cuda')
101
            else:
102
                raise RuntimeError('could not use CUDA on this machine')
103
        else:
104
            return torch.device('cpu')
105
106
    if isinstance(cuda, str):
107
        if 'cuda' in cuda:
108
            if torch.cuda.is_available():
109
                return torch.device(cuda)
110
            else:
111
                raise RuntimeError('could not use CUDA on this machine')
112
        elif 'cpu' in cuda:
113
            return torch.device('cpu')
114
        else:
115
            raise RuntimeError(
116
                'wrong device identifier'
117
                'see also: https://pytorch.org/docs/stable/tensor_attributes.html#torch.torch.device'
118
            )
119
120
    if isinstance(cuda, torch.device):
121
        return cuda