文章目录

一、先筛查

1. 检查当前默认设备（CPU/GPU）

2. 检查张量（Tensor）所在的设备

3. 检查模型（Model）所在的设备

二、后转换

1. 检查 CUDA 是否可用，如果可用，将模型移动到GPU（备注：具体模型定义变量参考你的代码）

2. 强制使用 CPU（即使 GPU 可用）

一、先筛查

1. 检查当前默认设备（CPU/GPU）

import torch

# 检查默认设备（返回当前默认的 torch.device 对象）
default_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Default device: {default_device}")  # 输出：cuda / cpu

2. 检查张量（Tensor）所在的设备

import torch

# 创建一个张量，默认在 CPU 上
x = torch.tensor([1, 2, 3])
print(x.device)  # 输出：cpu

3. 检查模型（Model）所在的设备

import torch
import torch.nn as nn

model = nn.Linear(10, 2)  # 定义一个简单的线性模型
print(next(model.parameters()).device)  # 输出：cpu（默认在 CPU 上）

二、后转换

1. 检查 CUDA 是否可用，如果可用，将模型移动到GPU（备注：具体模型定义变量参考你的代码）

import torch
import torch.nn as nn

# ①对象是张量——如果 CUDA 可用，手动将张量移动到 GPU
if torch.cuda.is_available():
    x = x.to('cuda')
    print(x.device)  # 输出：cuda:0（如果有多个 GPU，可能是 cuda:1 等）

# ②对象是模型——如果 CUDA 可用，将模型移动到 GPU
if torch.cuda.is_available():
    model = model.to('cuda')
    print(next(model.parameters()).device)  # 输出：cuda:0

2. 强制使用 CPU（即使 GPU 可用）

有时候，即使 GPU 可用，你可能希望强制代码在 CPU 上运行：

import torch

# 强制使用 CPU
device = torch.device('cpu')
x = torch.tensor([1, 2, 3], device=device)
model = model.to(device)