diff --git a/packages/plugin-legacy/index.js b/packages/plugin-legacy/index.js
index 0ecf2bd67b2ad5..60dee775aaa5d2 100644
--- a/packages/plugin-legacy/index.js
+++ b/packages/plugin-legacy/index.js
@@ -533,17 +533,13 @@ function isLegacyOutput(options) {
 function recordAndRemovePolyfillBabelPlugin(polyfills) {
   return ({ types: t }) => ({
     name: 'vite-remove-polyfill-import',
-    visitor: {
-      Program: {
-        exit(path) {
-          path.get('body').forEach((p) => {
-            if (t.isImportDeclaration(p)) {
-              polyfills.add(p.node.source.value)
-              p.remove()
-            }
-          })
+    post({ path }) {
+      path.get('body').forEach((p) => {
+        if (t.isImportDeclaration(p)) {
+          polyfills.add(p.node.source.value)
+          p.remove()
         }
-      }
+      })
     }
   })
 }