torch.cuda.set_device(local_rank)

这是 PyTorch 中设置使用哪个 GPU 设备的函数，其中 local_rank 表示当前进程所在的 GPU 的 ID。如果在多 GPU 环境下使用 PyTorch，需要将数据和模型放在不同的 GPU 上进行并行计算。使用该函数可以指定当前进程使用哪个 GPU 设备。例如，如果有 4 个 GPU 设备可用，可以分别将数据和模型放在 4 个不同的 GPU 设备上进行计算，从而实现并行计算加速模型训练。