diff --git a/ml-agents/mlagents/trainers/tests/mock_brain.py b/ml-agents/mlagents/trainers/tests/mock_brain.py
index 1fcfbbc710..75d7dfac33 100644
--- a/ml-agents/mlagents/trainers/tests/mock_brain.py
+++ b/ml-agents/mlagents/trainers/tests/mock_brain.py
@@ -129,7 +129,12 @@ def create_buffer(brain_infos, brain_params, sequence_length):
         buffer[0]["prev_action"].append(current_brain_info.previous_vector_actions[0])
         buffer[0]["masks"].append(1.0)
         buffer[0]["advantages"].append(1.0)
-        buffer[0]["action_probs"].append(np.ones(buffer[0]["actions"][0].shape))
+        if brain_params.vector_action_space_type == "discrete":
+            buffer[0]["action_probs"].append(
+                np.ones(sum(brain_params.vector_action_space_size))
+            )
+        else:
+            buffer[0]["action_probs"].append(np.ones(buffer[0]["actions"][0].shape))
         buffer[0]["actions_pre"].append(np.ones(buffer[0]["actions"][0].shape))
         buffer[0]["random_normal_epsilon"].append(
             np.ones(buffer[0]["actions"][0].shape)