diff options
| author | YurenHao0426 <Blackhao0426@gmail.com> | 2026-08-10 10:15:32 -0500 |
|---|---|---|
| committer | YurenHao0426 <Blackhao0426@gmail.com> | 2026-08-10 10:15:32 -0500 |
| commit | 20fd3577693593e3f62bcc4e1c2565dca5349944 (patch) | |
| tree | aa4b84f9267769982402dd07b34a69db52a9d6fc /experiments/babyai_shared_run.py | |
| parent | 6d99f411cf7e3f271c85976f27e3ad780f9ecbc8 (diff) | |
fix: resolve visible CUDA device for BabyAI
Diffstat (limited to 'experiments/babyai_shared_run.py')
| -rw-r--r-- | experiments/babyai_shared_run.py | 5 |
1 files changed, 4 insertions, 1 deletions
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 = ( |