La función torch.device de PyTorch


Pytorch torch 参考手册Manual de referencia de Pytorch torch

torch.deviceEs una función de PyTorch utilizada para crear objetos de dispositivo, que sirve para especificar el dispositivo de almacenamiento de tensores.

Definición de la función

torch.device(device)

Ejemplo de uso

Ejemplo

import torch

# Crear dispositivo CPU
cpu_device = torch.device('cpu')
print("Dispositivo CPU:", cpu_device)

# Crear dispositivo CUDA
if torch.cuda.is_available():
    cuda_device = torch.device('cuda:0')
    print("Dispositivo CUDA:", cuda_device)

    # Crear un tensor en la GPU
    x = torch.zeros(3, 4, device=cuda_device)
    print("Dispositivo del tensor:", x.device)

Pytorch torch 参考手册Manual de referencia de Pytorch torch

Otras extensiones