HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.3.15 new TPU device-selection ability (#16576)

Browse source 
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Ultralytics Assistant <135830346+UltralyticsAssistant@users.noreply.github.com>
Co-authored-by: Skillnoob <78843978+Skillnoob@users.noreply.github.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Kirill Lalayants 2024-10-17 03:57:01 +03:00 • committed by GitHub
parent ccda7ff973
commit 22ebd44f62
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 42 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

2
ultralytics/utils/torch_utils.py
View file

@ -163,7 +163,7 @@ def select_device(device="", batch=0, newline=False, verbose=True):
Note:
Sets the 'CUDA_VISIBLE_DEVICES' environment variable for specifying which GPUs to use.
"""
if isinstance(device, torch.device):
if isinstance(device, torch.device) or str(device).startswith("tpu"):
return device
s = f"Ultralytics {__version__} 🚀 Python-{PYTHON_VERSION} torch-{torch.__version__} "