Fix torch device option in argparser (#5467)

Unity-Technologies · Jul 23, 2021 · 5ffc783 · 5ffc783
1 parent 6d8ba75
commit 5ffc783
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 2 deletions.
diff --git a/com.unity.ml-agents/CHANGELOG.md b/com.unity.ml-agents/CHANGELOG.md
@@ -18,6 +18,7 @@ and this project adheres to
 #### ml-agents / ml-agents-envs / gym-unity (Python)
 - Fixed a bug in multi-agent cooperative training where agents might not receive all of the states of
 terminated teammates. (#5441)
+- Fixed wrong attribute name in argparser for torch device option (#5433)(#5467)
 
 ## [2.1.0-exp.1] - 2021-06-09
 ### Minor Changes

diff --git a/ml-agents/mlagents/trainers/cli_utils.py b/ml-agents/mlagents/trainers/cli_utils.py
@@ -257,7 +257,7 @@ def _create_parser() -> argparse.ArgumentParser:
     torch_conf.add_argument(
         "--torch-device",
         default=None,
-        dest="torch_device",
+        dest="device",
         action=DetectDefault,
         help='Settings for the default torch.device used in training, for example, "cpu", "cuda", or "cuda:0"',
     )

diff --git a/ml-agents/mlagents/trainers/settings.py b/ml-agents/mlagents/trainers/settings.py
@@ -797,7 +797,7 @@ class EngineSettings:
 
 @attr.s(auto_attribs=True)
 class TorchSettings:
-    device: Optional[str] = parser.get_default("torch_device")
+    device: Optional[str] = parser.get_default("device")
 
 
 @attr.s(auto_attribs=True)