From 1c3d3cf6cfc4c3e69326e414303f6952394cd776 Mon Sep 17 00:00:00 2001
From: Dominik Roth <dominik.roth.dev@gmail.com>
Date: Sat, 9 Mar 2024 13:58:56 +0100
Subject: [PATCH] Fix typo 2

---
 metastable_baselines2/common/distributions.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/metastable_baselines2/common/distributions.py b/metastable_baselines2/common/distributions.py
index ab69dd4..fd80e06 100644
--- a/metastable_baselines2/common/distributions.py
+++ b/metastable_baselines2/common/distributions.py
@@ -24,9 +24,10 @@ def _patched_make_proba_distribution(
             cls = StateDependentNoiseDistribution
         elif use_pca:
             cls = PCA_Distribution
+            return cls(get_action_dim(action_space), n_envs=n_envs, **dist_kwargs)
         else:
             cls = DiagGaussianDistribution
-        return cls(get_action_dim(action_space), n_envs=n_envs, **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):