From 550762390a86ec1e1d7f250c034da8db44deaeed Mon Sep 17 00:00:00 2001 From: clemente0731 Date: Wed, 3 Jul 2024 10:51:11 +0800 Subject: [PATCH] fix device issue when host don't have cuda device --- .../xtest_device_conversion_detection.py | 14 +++++---- pprobe/toggle/cli.py | 29 ++++++++++++++----- 2 files changed, 31 insertions(+), 12 deletions(-) diff --git a/pprobe/tests/xtest_device_conversion_detection.py b/pprobe/tests/xtest_device_conversion_detection.py index e4424d5..f80ef52 100644 --- a/pprobe/tests/xtest_device_conversion_detection.py +++ b/pprobe/tests/xtest_device_conversion_detection.py @@ -1,12 +1,13 @@ import torch import torch.nn as nn -# Create a tensor on GPU -tensor = torch.randn(3, 3).cuda() +# 获取可用设备 +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -# Create a model on GPU -model = nn.Linear(3, 3).cuda() +tensor = torch.randn(3, 3).to(device) +model = nn.Linear(3, 3).to(device) +# 将张量移动到 CPU 的三种方法 # Method 1: Using tensor.to("cpu") tensor_to_cpu_1 = tensor.to("cpu") @@ -16,6 +17,7 @@ # Method 3: Using tensor.cpu() tensor_to_cpu_3 = tensor.cpu() +# 将模型移动到 CPU 的三种方法 # Method 1: Using model.to("cpu") model_to_cpu_1 = model.to("cpu") @@ -23,4 +25,6 @@ model_to_cpu_2 = model.to(torch.device("cpu")) # Method 3: Using model.cpu() -model_to_cpu_3 = model.cpu() \ No newline at end of file +model_to_cpu_3 = model.cpu() + +print(f"Using device: {device}") \ No newline at end of file diff --git a/pprobe/toggle/cli.py b/pprobe/toggle/cli.py index 36be75a..874d876 100644 --- a/pprobe/toggle/cli.py +++ b/pprobe/toggle/cli.py @@ -20,13 +20,28 @@ ================================================= ██████╗ ██████╗ ██████╗ ██████╗ ██████╗ ███████╗ ██╔══██╗██╔══██╗██╔══██╗██╔═══██╗██╔══██╗██╔════╝ -██████╔╝██████╔╝██████╔╝██║ ██║██████╔╝█████╗ -██╔═══╝ ██╔═══╝ ██╔══██╗██║ ██║██╔══██╗██╔══╝ +██████╔╝██████╔╝██████╔╝██║ ██║██████╔╝█████╗ +██╔═══╝ ██╔═══╝ ██╔══██╗██║ ██║██╔══██╗██╔══╝ ██║ ██║ ██║ ██║╚██████╔╝██████╔╝███████╗ -╚═╝ ╚═╝ ╚═╝ ╚═╝ ╚═════╝ ╚═════╝ ╚══════╝ +╚═╝ ╚═╝ ╚═╝ ╚═╝ ╚═════╝ ╚═════╝ ╚══════╝ ================================================= MIT License -Copyright (c) [2024] [clemente0620@gmail.com] + +Copyright (c) 2024 clemente0620@gmail.com + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. """ @@ -131,9 +146,9 @@ def show_status(self): True values in green False values in red """ - + print(HELLO_PPROBE) - + def colorize(value): """Return the value in green if True, in red if False.""" if value in (True, "True"): @@ -141,7 +156,7 @@ def colorize(value): elif value in (False, "False"): return f"\033[91m{value}\033[0m" # RED return value - + status_in_color = [] for name, value in self.running_toggle.items():