diff --git a/torchy_baselines/common/buffers.py b/torchy_baselines/common/buffers.py
index 565117c..b169dd3 100644
--- a/torchy_baselines/common/buffers.py
+++ b/torchy_baselines/common/buffers.py
@@ -134,6 +134,7 @@ class RolloutBuffer(BaseBuffer):
         if len(log_prob.shape) == 0:
             # Reshape 0-d tensor to avoid error
             log_prob = log_prob.reshape(-1, 1)
+
         self.observations[self.pos] = th.FloatTensor(np.array(obs).copy())
         self.actions[self.pos] = th.FloatTensor(np.array(action).copy())
         self.rewards[self.pos] = th.FloatTensor(np.array(reward).copy())
diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py
index 3d748fb..b6584dd 100644
--- a/torchy_baselines/ppo/ppo.py
+++ b/torchy_baselines/ppo/ppo.py
@@ -158,6 +158,9 @@ class PPO(BaseRLModel):
 
             self._update_info_buffer(infos)
             n_steps += 1
+            if isinstance(self.action_space, gym.spaces.Discrete):
+                # Reshape in case of discrete action
+                actions = actions.reshape(-1, 1)
             rollout_buffer.add(obs, actions, rewards, dones, values, log_probs)
             obs = new_obs