在 PyTorch 中，你可以使用 CUDA_VISIBLE_DEVICES 环境变量来指定程序使用的 GPU。在你的命令行中运行脚本之前，你可以设置该环境变量。以下是如何指定显卡的一种方法：

假设使用第一张 GPU，可以运行以下命令：

CUDA_VISIBLE_DEVICES=0 python main.py

这将使程序只能看到第一张 GPU，并使用它进行计算。如果想使用多个 GPU，你可以在等号后面提供逗号分隔的 GPU 索引列表：

CUDA_VISIBLE_DEVICES=0,1 python main.py

在这个例子中，程序将只能看到索引为 0 和 1 的两张 GPU，并使用它们进行计算。

        记得在运行之前确保 PyTorch 可以使用 CUDA，即 torch.cuda.is_available() 返回 True。在你的代码中，你已经使用了 args.cuda = torch.cuda.is_available() 来设置这个变量，确保 CUDA 可用性。

        请注意，使用这种方法只是在运行脚本时指定程序能够看到哪些 GPU，而不是在代码中直接指定使用某一张 GPU。如果想在代码中指定使用特定的 GPU，可以在代码中添加相应的 PyTorch CUDA 设置，例如：

# 在你的代码中添加以下部分
if args.cuda:torch.cuda.set_device(0)  # 设置使用第一张 GPU# 其他 CUDA 相关设置

        这将在代码中明确指定程序使用的 GPU 设备。根据你的需求，你可以在 set_device 中传递不同的 GPU 索引。