diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py
index b0752e2..ce7e699 100644
--- a/torchy_baselines/common/distributions.py
+++ b/torchy_baselines/common/distributions.py
@@ -288,7 +288,7 @@ class StateDependentNoiseDistribution(Distribution):
         if self.full_std:
             return std
         # Reduce the number of parameters:
-        return th.ones((self.latent_sde_dim, self.action_dim)).to(log_std.device) * std
+        return th.ones(self.latent_sde_dim, self.action_dim).to(log_std.device) * std
 
     def sample_weights(self, log_std):
         """
diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py
index 34bb3e4..0774ef0 100644
--- a/torchy_baselines/sac/policies.py
+++ b/torchy_baselines/sac/policies.py
@@ -33,7 +33,7 @@ class Actor(BaseNetwork):
         if self.use_sde:
             # TODO: check for the learn_features
             self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
-                                                               learn_features=False, squash_output=True)
+                                                               learn_features=True, squash_output=True)
             self.mu, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1],
                                                                             log_std_init=log_std_init)
         else:
diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py
index 49f562f..5c763b5 100644
--- a/torchy_baselines/sac/sac.py
+++ b/torchy_baselines/sac/sac.py
@@ -170,11 +170,13 @@ class SAC(BaseRLModel):
 
             obs, action_batch, next_obs, done, reward = replay_data
 
-            # TODO: check if there is another way to fix pytorch complain
-            # if we don't sample the weights again
-            # (Trying to backward through the graph a second time)
-            if self.use_sde:
-                self.actor.reset_noise()
+            # Two options: retain_graph=True in the actor_loss.backward()
+            # or sample again the noise matrix
+            # otherwise the intermediate step `std = th.exp(log_std)`
+            # is lost and we cannot backpropagate through again
+            # if self.use_sde:
+            #     self.actor.reset_noise()
+
             # Action by the current actor for the sampled state
             action_pi, log_prob = self.actor.action_log_prob(obs)
             log_prob = log_prob.reshape(-1, 1)
@@ -196,10 +198,10 @@ class SAC(BaseRLModel):
                 ent_coef_loss.backward()
                 self.ent_coef_optimizer.step()
 
-            # Select action according to policy
-            next_action, next_log_prob = self.actor.action_log_prob(next_obs)
 
             with th.no_grad():
+                # Select action according to policy
+                next_action, next_log_prob = self.actor.action_log_prob(next_obs)
                 # Compute the target Q value
                 target_q1, target_q2 = self.critic_target(next_obs, next_action)
                 target_q = th.min(target_q1, target_q2)
@@ -227,7 +229,10 @@ class SAC(BaseRLModel):
 
             # Optimize the actor
             self.actor.optimizer.zero_grad()
-            actor_loss.backward()
+            # Cf comment above, otherwise pytorch raises an error
+            # ("Trying to backward through the graph a second time")
+            retain_graph = True if self.use_sde and gradient_steps > 1 else False
+            actor_loss.backward(retain_graph=retain_graph)
             self.actor.optimizer.step()
 
             # Update target networks
@@ -239,6 +244,8 @@ class SAC(BaseRLModel):
         logger.logkv("ent_coef", ent_coef.item())
         logger.logkv("actor_loss", actor_loss.item())
         logger.logkv("critic_loss", critic_loss.item())
+        if ent_coef_loss is not None:
+            logger.logkv("ent_coef_loss", ent_coef_loss.item())
 
     def learn(self, total_timesteps, callback=None, log_interval=4,
               eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="SAC",