Revision history [back]

The recommended approach in the most recent release of PyTorch is to use the torch.device and to() methods to specify and transfer tensors to the desired device.

First, define the device using torch.device, for example:

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

Then, move the tensor to this device using the to() method:

tensor = tensor.to(device)

This will automatically transfer the tensor to the GPU (if available) or CPU, depending on the device variable.

To ensure that all the tensors in your model are on the same device, you can define the device at the beginning of your code and use the to() method to transfer tensors to the device as needed:

# Define device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# Define model
model = MyModel().to(device)

# Load data
data = DataLoader(...)

# Train model
for batch in data:
  inputs, labels = batch[0].to(device), batch[1].to(device)
  predictions = model(inputs)
  ...

This will ensure that all the tensors in the model (including the model parameters) are automatically transferred to the desired device.