From 5c490f1b8540af5fe08486b0eceb004ae4578518 Mon Sep 17 00:00:00 2001 From: robinDLFM Date: Tue, 2 Sep 2025 12:03:06 +0200 Subject: [PATCH] Fix: data on both cuda:0 and cuda:1 when using --torch-device cuda:1 --- ml-agents/mlagents/torch_utils/torch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ml-agents/mlagents/torch_utils/torch.py b/ml-agents/mlagents/torch_utils/torch.py index 311304ef54..cf819fb807 100644 --- a/ml-agents/mlagents/torch_utils/torch.py +++ b/ml-agents/mlagents/torch_utils/torch.py @@ -52,7 +52,7 @@ def set_torch_config(torch_settings: TorchSettings) -> None: _device = torch.device(device_str) if _device.type == "cuda": - torch.set_default_device(_device.type) + torch.set_default_device(_device) torch.set_default_dtype(torch.float32) else: torch.set_default_dtype(torch.float32)