From 1660b5275136acc63f992a1f5f70ef6d2b00ad4b Mon Sep 17 00:00:00 2001
From: Eric Liang <ekhliang@gmail.com>
Date: Sun, 23 Feb 2020 13:13:43 -0800
Subject: [PATCH] [rllib] Fix torch GPU / yaml load warning (#7278)

* fix

* safe load

* reduce num buffer shardscZZ
---
 rllib/tests/test_eager_support.py | 5 ++++-
 rllib/utils/from_config.py        | 4 ++--
 rllib/utils/torch_ops.py          | 2 +-
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/rllib/tests/test_eager_support.py b/rllib/tests/test_eager_support.py
index 0e36557fe..6a02f6355 100644
--- a/rllib/tests/test_eager_support.py
+++ b/rllib/tests/test_eager_support.py
@@ -60,7 +60,10 @@ class TestEagerSupport(unittest.TestCase):
                 "learning_starts": 0,
                 "num_gpus": 0,
                 "min_iter_time_s": 1,
-                "timesteps_per_iteration": 100
+                "timesteps_per_iteration": 100,
+                "optimizer": {
+                    "num_replay_buffer_shards": 1,
+                },
             })
 
     def testSAC(self):
diff --git a/rllib/utils/from_config.py b/rllib/utils/from_config.py
index 09073891a..dc401284d 100644
--- a/rllib/utils/from_config.py
+++ b/rllib/utils/from_config.py
@@ -131,7 +131,7 @@ def from_config(cls, config=None, **kwargs):
             if re.search("\\.(yaml|yml|json)$", type_):
                 return from_file(cls, type_, *ctor_args, **ctor_kwargs)
             # Try un-json/un-yaml'ing the string into a dict.
-            obj = yaml.load(type_)
+            obj = yaml.safe_load(type_)
             if isinstance(obj, dict):
                 return from_config(cls, obj)
             try:
@@ -208,7 +208,7 @@ def from_file(cls, filename, *args, **kwargs):
 
     with open(path, "rt") as fp:
         if path.endswith(".yaml") or path.endswith(".yml"):
-            config = yaml.load(fp)
+            config = yaml.safe_load(fp)
         else:
             config = json.load(fp)
 
diff --git a/rllib/utils/torch_ops.py b/rllib/utils/torch_ops.py
index 799ec09dc..c762b0670 100644
--- a/rllib/utils/torch_ops.py
+++ b/rllib/utils/torch_ops.py
@@ -34,7 +34,7 @@ def convert_to_non_torch_type(stats_dict):
     ret = {}
     for k, v in stats_dict.items():
         if isinstance(v, torch.Tensor):
-            ret[k] = v.item() if len(v.size()) == 0 else v.numpy()
+            ret[k] = v.cpu().item() if len(v.size()) == 0 else v.cpu().numpy()
         else:
             ret[k] = v
     return ret