安装 PyTorch CUDA

conda deactivate
conda env remove -n pytorch_env -y
rm -rf ~/miniconda3/envs/pytorch_env  # 确保彻底删除残留文件

conda create -n pytorch_env python=3.10 -y
conda activate pytorch_env

# 使用官方推荐命令安装（匹配CUDA 12.1）
conda install -c pytorch -c conda-forge \
  pytorch=2.5.1 \
  torchvision=0.20.1 \
  torchaudio=2.5.1 \
  pytorch-cuda=12.1 \
  --strict-channel-priority

python -c "import torch; import torchvision; \
print(f'Torch: {torch.__version__}\\nVision: {torchvision.__version__}\\nCUDA: {torch.cuda.is_available()}'); \
print(torch.__config__.show())"

# Verify the installation
import torch
print(torch.__version__)  # E.g., 2.4.0
print(torch.cuda.is_available())  # Should return True
print(torch.version.cuda)  # Should show 12.1
print(torch.cuda.get_device_name(0))  # E.g., NVIDIA GeForce RTX 3080
print(torch.backends.cudnn.enabled)  # Should return True
print(torch.backends.cudnn.version())  # E.g., 90000 for cuDNN 9.x
print(torch.cuda.nccl.version())  # Should show the NCCL version

import torch

x = torch.rand(5, 3).cuda()  # Create a random tensor on the GPU
y = torch.rand(3, 3).cuda()  # Another random tensor on the GPU
result = torch.matmul(x, y)

print(torch.cuda.is_available())
print(result)
print("Tensor is on GPU:", result.is_cuda)

True
tensor([[0.7633, 0.4071, 0.7316],
        [0.6251, 0.3729, 0.5449],
        [0.9465, 0.4374, 0.9828],
        [1.1408, 0.5835, 1.0626],
        [0.9263, 0.5343, 0.7289]], device='cuda:0')
Tensor is on GPU: True