1 files changed, 2 insertions, 2 deletions
diff --git a/app/assets/javascripts/image_diff/helpers/dom_helper.js b/app/assets/javascripts/image_diff/helpers/dom_helper.js
index a61e5f01f9b..3468a629f5a 100644
--- a/app/assets/javascripts/image_diff/helpers/dom_helper.js
+++ b/app/assets/javascripts/image_diff/helpers/dom_helper.js
@@ -10,12 +10,12 @@ export function setPositionDataAttribute(el, options) {
 }
 
 export function updateDiscussionAvatarBadgeNumber(discussionEl, newBadgeNumber) {
-  const avatarBadgeEl = discussionEl.querySelector('.image-diff-avatar-link .badge');
+  const avatarBadgeEl = discussionEl.querySelector('.image-diff-avatar-link .design-note-pin');
   avatarBadgeEl.textContent = newBadgeNumber;
 }
 
 export function updateDiscussionBadgeNumber(discussionEl, newBadgeNumber) {
-  const discussionBadgeEl = discussionEl.querySelector('.badge');
+  const discussionBadgeEl = discussionEl.querySelector('.design-note-pin');
   discussionBadgeEl.textContent = newBadgeNumber;
 }