diff --git a/public/embed-example.html b/public/embed-example.html
index 14379a2d..8345738d 100644
--- a/public/embed-example.html
+++ b/public/embed-example.html
@@ -27,7 +27,7 @@ <h3>Both search and initial questions, opening inline, expanding iframe height,
     const iframes = [...document.querySelectorAll('iframe')]
     window.addEventListener('message', function (event) {
       if (event.data.type === 'aisafety.info__height') {
-        const iframe = iframes.find((iframe) => iframe.src === event.source.location.href)
+        const iframe = event.source.frameElement
         if (iframe) {
           iframe.style.height = event.data.height + 'px'
         }