Fixed discrete models

ml-agents/mlagents/trainers/torch/model_serialization.py

             + ["action_masks", "memories"]
         )
 
+        if self.policy.use_continuous_act:
+            action_name = "action"
+            action_prob_name = "action_probs"
+        else:
+            action_name = "action_unused"
+            action_prob_name = "action"
+
-            "action",
-            "action_probs",
+            action_name,
+            action_prob_name,
             "version_number",
             "memory_size",
             "is_continuous_control",