saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-23 19:32:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Revert "Do not copy and try channel last format"

Browse source 
This reverts commit 1103481c26.
This commit is contained in:
Antonin Raffin 2022-12-17 12:41:43 +01:00
parent 83f2068274
commit 58b1e74abb
No known key found for this signature in database
GPG key ID: B8B48F65CAD6232C
4 changed files with 7 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
stable_baselines3/common/buffers.py
View file

@ -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(

2
stable_baselines3/common/off_policy_algorithm.py
View file

@ -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()

2
stable_baselines3/common/on_policy_algorithm.py
View file

@ -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,

7
stable_baselines3/common/utils.py
View file

@ -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)}")