怎样在 PyTorch 中选择使用 CPU 还是 GPU？

在 PyTorch 中，你可以通过以下几个步骤来选择使用 CPU 或 GPU：

1. **检查可用设备**：首先，可以使用 `torch.cuda.is_available()` 函数检查当前设备是否支持 GPU。如果返回 `True`，说明有可用的 GPU；反之则是 CPU。

import torch
if torch.cuda.is_available():
    device = torch.device("cuda")  # 如果有GPU，则设为GPU模式
else:
    device = torch.device("cpu")  # 否则默认为CPU

2. **明确指定设备**：当你创建张量（tensor）时，可以直接指定其所在的设备。例如，如果你已经设置了 `device` 为 GPU，创建张量时就用 `.to(device)`：

x = torch.randn(5, 5)
x_gpu = x.to(device)  # 将x移动到指定设备

3. **自动选择**：PyTorch 提供了自动选择策略，如 `.to('cuda:0')`，它会在第一个可用的 GPU 设备上运行，如果没有找到 GPU，则自动切换到 CPU。

4. **模型和优化器**：在构建模型和优化器时，也可以设置它们使用的设备。例如：

model = YourModel().to(device)
optimizer = torch.optim.Adam(model.parameters(), device=device)