diff --git a/mindmap/mindmap.py b/mindmap/mindmap.py
index 63fea55..ceb7608 100644
--- a/mindmap/mindmap.py
+++ b/mindmap/mindmap.py
@@ -518,7 +518,7 @@ def update_student_submission_status(self, student_id, submission_status):
         """
         from submissions.api import create_submission
         student_submission = self.get_submission(student_id)
-        mindmap_student_body = student_submission.get("mindmap_student_body", {})
+        mindmap_student_body = student_submission.get("answer", {}).get("mindmap_student_body", {})
         answer = {
             "mindmap_student_body": mindmap_student_body,
             "submission_status": submission_status,
@@ -539,20 +539,21 @@ def remove_grade(self, data, _suffix="") -> dict:
             dict: A dictionary containing the handler result.
         """
         # Lazy import: import here to avoid app not ready errors
-        from submissions.api import reset_score, create_submission # pylint: disable=import-outside-toplevel
+        from submissions.api import reset_score # pylint: disable=import-outside-toplevel
 
         require(self.is_instructor())
 
         student_id = data.get("student_id")
-        if not student_id:
+        submission_id = data.get("submission_id")
+        if not student_id or not submission_id:
             raise JsonHandlerError(400, "Missing required parameters")
 
-        latest_submission = self.update_student_submission_status(
+        reset_score(submission_id, self.block_course_id, self.block_id)
+
+        self.update_student_submission_status(
             data.get("student_id"), SubmissionStatus.SUBMITTED.value,
         )
 
-        reset_score(latest_submission["uuid"], self.block_course_id, self.block_id)
-
         return {
             "success": True,
         }
diff --git a/mindmap/public/js/src/mindmap.js b/mindmap/public/js/src/mindmap.js
index 108211d..09980ee 100644
--- a/mindmap/public/js/src/mindmap.js
+++ b/mindmap/public/js/src/mindmap.js
@@ -227,7 +227,7 @@ function MindMapXBlock(runtime, element, context) {
                     e.preventDefault();
                     const typeAction = $(this).attr("data-type");
                     const grade = $("#grade_value").val();
-                    const { student_id } = submissionData;
+                    const { student_id, submission_id } = submissionData;
                     const invalidGradeMessage = gettext("Invalid grade must be a number");
                     const maxGradeMessage = gettext("Please enter a lower grade, maximum grade allowed is:");
                     const gradeParsed = parseInt(grade, 10);
@@ -261,6 +261,7 @@ function MindMapXBlock(runtime, element, context) {
                       apiUrl = removeGradeURL;
                       data = {
                         student_id: student_id,
+                        submission_id: submission_id,
                       };
                     }