diff options
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 = ( |