How to use 'torch check gpu' in Python

Every line of 'torch check gpu' code snippets is scanned for vulnerabilities by our powerful machine learning engine that combs millions of open source libraries, ensuring your Python code is secure.

All examples are scanned by Snyk Code

By copying the Snyk Code Snippets you agree to
this disclaimer
95@staticmethod
96def 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
Important

Use secure code every time

Secure your code as it's written. Use Snyk Code to scan source code in minutes – no build needed – and fix issues immediately. Enable Snyk Code


Related snippets