From 83f206827476d89b6c41d1a744d1002a9bee9236 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Sat, 17 Dec 2022 12:41:28 +0100 Subject: [PATCH] Revert "Change memory format only for images" This reverts commit 864b0d9d7f62269c36f4372680e24d9156518a61. --- stable_baselines3/common/buffers.py | 5 ++--- stable_baselines3/common/utils.py | 8 +++++--- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index 768f51e..16f206f 100644 --- a/stable_baselines3/common/buffers.py +++ b/stable_baselines3/common/buffers.py @@ -131,10 +131,9 @@ class BaseBuffer(ABC): (may be useful to avoid changing things be reference) :return: """ - kwargs = {"memory_format": th.channels_last} if len(array.shape) >= 4 else {} if copy: - return th.tensor(array).to(self.device, non_blocking=True, **kwargs) - return th.as_tensor(array).to(self.device, non_blocking=True, **kwargs) + return th.tensor(array).to(self.device, memory_format=th.channels_last, non_blocking=True) + return th.as_tensor(array).to(self.device, memory_format=th.channels_last, non_blocking=True) @staticmethod def _normalize_obs( diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py index 777d146..b713a28 100644 --- a/stable_baselines3/common/utils.py +++ b/stable_baselines3/common/utils.py @@ -460,11 +460,13 @@ def obs_as_tensor( :param device: PyTorch device :return: PyTorch tensor of the observation on a desired device. """ - kwargs = {"memory_format": th.channels_last} if len(obs.shape) >= 4 else {} if isinstance(obs, np.ndarray): - return th.as_tensor(obs).to(device, non_blocking=True, **kwargs) + return th.as_tensor(obs).to(device, memory_format=th.channels_last, non_blocking=True) elif isinstance(obs, dict): - return {key: th.as_tensor(_obs).to(device, non_blocking=True, **kwargs) for (key, _obs) in obs.items()} + return { + key: th.as_tensor(_obs).to(device, memory_format=th.channels_last, non_blocking=True) + for (key, _obs) in obs.items() + } else: raise Exception(f"Unrecognized type of observation {type(obs)}")