torch.cuda.set_device(): Selectarea GPU-ului Activ în Medii Multi-GPU în PyTorch & Deep Learning Practic

Funcția torch.cuda.set_device() stabilește indexul plăcii grafice curente (de la 0 la device_count() - 1) pentru procesul de execuție, direcționând automat instrucțiunile .cuda() către placa specificată.

Exemplu pe mașini cu 4 sau 8 GPU-uri

import torch

if torch.cuda.is_available() and torch.cuda.device_count() > 1:
    # Selectarea celui de-al doilea GPU (index 1)
    torch.cuda.set_device(1)
    x = torch.randn(10, 10).cuda()
    print("Tensorul este alocat pe:", x.device)  # cuda:1

În antrenarea distribuită este frecvent utilizată împreună cu variabila de mediu LOCAL_RANK.