From 20fd3577693593e3f62bcc4e1c2565dca5349944 Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Mon, 10 Aug 2026 10:15:32 -0500 Subject: fix: resolve visible CUDA device for BabyAI --- experiments/babyai_shared_run.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) (limited to 'experiments') diff --git a/experiments/babyai_shared_run.py b/experiments/babyai_shared_run.py index 1b460b7..6b6ddd2 100644 --- a/experiments/babyai_shared_run.py +++ b/experiments/babyai_shared_run.py @@ -129,7 +129,10 @@ def main(): if device.type == "cpu": torch.set_num_threads(1) if device.type == "cuda": - torch.cuda.set_device(device) + device_index = (torch.cuda.current_device() + if device.index is None else device.index) + torch.cuda.set_device(device_index) + device = torch.device("cuda", device_index) torch.cuda.reset_peak_memory_stats(device) metadata, train, validation, rollout_seeds = load_data(args.data) cardinalities = ( -- cgit v1.2.3