1.1 Approach 1: Using os.environ['CUDA_VISIBLE_DEVICES']

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '2'
model = NeuralNet().cuda()
batch = batch.cuda()

1.2 Approach 2: Using torch.device()

target_device = torch.device('cuda:2')
model = NeuralNet().to(target_device)
batch = batch.to(target_device)

1.3 Errer 1: RuntimeError: CUDA error: invalid device ordinal

Solution: Do not mix Approach 1 and Approach 2 in the same script.

1.4 torch.load Error: RuntimeError: CUDA out of memory

Solution:

1. If GPU memory is insufficient, you must switch to a different GPU.

2. Even when using Approach 2 to specify a device, torch.load loads from GPU 0 by default (the card where the model was saved).

3. Use Approach 1 to specify the GPU, or use torch.load(path, map_location=lambda storage, loc: storage.cuda(2)) too keep loaded tensors on the specified GPU.

4. Alternatively, use torch.load(path, map_location='cpu') to load everything on to CPU.

5. Selecting Specific GPUs (Multi-GPU with DataParallel)

2.1 Standard DataParallel Usage

DataParallel is straightforward to use—simply wrap the model with DataParallel. Training code remains the same as single-GPU training.

import torch
from torch.nn import DataParallel

network = CustomModel()
network = DataParallel(network, device_ids=[0, 1])

2.2 DataParallel Usage in PyTorch Geometric

Important: When using torch_geometric (PyG), you must use DataListLoader from torch_geometric.nn!

gpu_list = [0, 2, 3]
# Must specify the primary device (defaults to GPU 0)
# Not specifying the device causes: RuntimeError: Expected all tensors to be on the same device
primary_device = torch.device(f'cuda:{gpu_list[0]}')

network = GraphNet()
network = DataParallel(network, device_ids=gpu_list)
network.to(primary_device)

training_data = GraphDataset(params)
training_loader = DataListLoader(training_data, batch_size=batch_size, shuffle=True)

def train_epoch(net, dataloader):
    net.train()
    for idx, graph_batch in enumerate(dataloader):
        predictions = net(graph_batch).float()
        labels = torch.tensor([g.y for g in graph_batch]).to(predictions.device)