From 58b1e74abbbc7df16168151c5b8f690cafda3671 Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Sat, 17 Dec 2022 12:41:43 +0100 Subject: [PATCH] Revert "Do not copy and try channel last format" This reverts commit 1103481c26f6de0e13abbd1121ad1789958d84e5. --- stable_baselines3/common/buffers.py | 6 +++--- stable_baselines3/common/off_policy_algorithm.py | 2 +- stable_baselines3/common/on_policy_algorithm.py | 2 +- stable_baselines3/common/utils.py | 7 ++----- 4 files changed, 7 insertions(+), 10 deletions(-) diff --git a/stable_baselines3/common/buffers.py b/stable_baselines3/common/buffers.py index 16f206f..09b80c3 100644 --- a/stable_baselines3/common/buffers.py +++ b/stable_baselines3/common/buffers.py @@ -121,7 +121,7 @@ class BaseBuffer(ABC): """ raise NotImplementedError() - def to_torch(self, array: np.ndarray, copy: bool = False) -> th.Tensor: + def to_torch(self, array: np.ndarray, copy: bool = True) -> th.Tensor: """ Convert a numpy array to a PyTorch tensor. Note: it copies the data by default @@ -132,8 +132,8 @@ class BaseBuffer(ABC): :return: """ if copy: - 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) + return th.tensor(array).to(self.device, non_blocking=True) + return th.as_tensor(array).to(self.device, non_blocking=True) @staticmethod def _normalize_obs( diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index 20b838e..5e018fc 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/stable_baselines3/common/off_policy_algorithm.py @@ -218,7 +218,7 @@ class OffPolicyAlgorithm(BaseAlgorithm): self.lr_schedule, **self.policy_kwargs, # pytype:disable=not-instantiable ) - self.policy = self.policy.to(self.device, memory_format=th.channels_last, non_blocking=True) + self.policy = self.policy.to(self.device) # Convert train freq parameter to TrainFreq object self._convert_train_freq() diff --git a/stable_baselines3/common/on_policy_algorithm.py b/stable_baselines3/common/on_policy_algorithm.py index 8ce75f9..35ad2b9 100644 --- a/stable_baselines3/common/on_policy_algorithm.py +++ b/stable_baselines3/common/on_policy_algorithm.py @@ -121,7 +121,7 @@ class OnPolicyAlgorithm(BaseAlgorithm): use_sde=self.use_sde, **self.policy_kwargs # pytype:disable=not-instantiable ) - self.policy = self.policy.to(self.device, memory_format=th.channels_last, non_blocking=True) + self.policy = self.policy.to(self.device) def collect_rollouts( self, diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py index b713a28..b7c4cb9 100644 --- a/stable_baselines3/common/utils.py +++ b/stable_baselines3/common/utils.py @@ -461,12 +461,9 @@ def obs_as_tensor( :return: PyTorch tensor of the observation on a desired device. """ if isinstance(obs, np.ndarray): - return th.as_tensor(obs).to(device, memory_format=th.channels_last, non_blocking=True) + return th.as_tensor(obs).to(device, non_blocking=True) elif isinstance(obs, dict): - return { - key: th.as_tensor(_obs).to(device, memory_format=th.channels_last, non_blocking=True) - for (key, _obs) in obs.items() - } + return {key: th.as_tensor(_obs).to(device, non_blocking=True) for (key, _obs) in obs.items()} else: raise Exception(f"Unrecognized type of observation {type(obs)}")