From cc7a58bc5f3ad3e4e610a54711afa9246324adee Mon Sep 17 00:00:00 2001 From: Adam Gleave Date: Thu, 2 Jul 2020 20:28:48 -0700 Subject: [PATCH] Bugfix --- stable_baselines3/common/distributions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/stable_baselines3/common/distributions.py b/stable_baselines3/common/distributions.py index 068bba5..6fa70f8 100644 --- a/stable_baselines3/common/distributions.py +++ b/stable_baselines3/common/distributions.py @@ -655,7 +655,7 @@ def make_proba_distribution(action_space: gym.spaces.Space, if isinstance(action_space, spaces.Box): assert len(action_space.shape) == 1, "Error: the action space must be a vector" cls = StateDependentNoiseDistribution if use_sde else DiagGaussianDistribution - cls(get_action_dim(action_space), **dist_kwargs) + return cls(get_action_dim(action_space), **dist_kwargs) elif isinstance(action_space, spaces.Discrete): return CategoricalDistribution(action_space.n, **dist_kwargs) elif isinstance(action_space, spaces.MultiDiscrete):