add support for mac m1/m2 chips mps gpu

need upgrade onnx and protobuf, add setup default device for torch, so user can use command params "--torch-device mps:0" to launch gpu training on mac
Unity-Technologies · Mar 31, 2024 · 1853609 · 1853609
1 parent fb2af76
commit 1853609
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/ml-agents/mlagents/torch_utils/torch.py b/ml-agents/mlagents/torch_utils/torch.py
@@ -53,6 +53,8 @@ def set_torch_config(torch_settings: TorchSettings) -> None:
 
     if _device.type == "cuda":
         torch.set_default_tensor_type(torch.cuda.FloatTensor)
+    elif _device.type == 'mps':
+        torch.set_default_device(device_str)
     else:
         torch.set_default_tensor_type(torch.FloatTensor)
     logger.debug(f"default Torch device: {_device}")