diff --git a/Sources/Extensions/UIViewController+Hero.swift b/Sources/Extensions/UIViewController+Hero.swift
index cc4527f1..8c0a270c 100644
--- a/Sources/Extensions/UIViewController+Hero.swift
+++ b/Sources/Extensions/UIViewController+Hero.swift
@@ -336,9 +336,7 @@ public extension HeroExtension where Base: UIViewController {
             parentVC.present(next, animated: false, completion: completion)
           }
         } else {
-          #if TARGET_IS_EXTENSION
-          UIApplication.shared.keyWindow?.rootViewController = next
-          #endif
+          parentVC?.view.window?.rootViewController = next
         }
       }
     }
diff --git a/Sources/Transition/HeroTransition+Complete.swift b/Sources/Transition/HeroTransition+Complete.swift
index 208b97b5..f4b87a61 100644
--- a/Sources/Transition/HeroTransition+Complete.swift
+++ b/Sources/Transition/HeroTransition+Complete.swift
@@ -89,9 +89,7 @@ extension HeroTransition {
       if isPresenting != finished, !inContainerController, transitionContext != nil {
         // only happens when present a .overFullScreen VC
         // bug: http://openradar.appspot.com/radar?id=5320103646199808
-        #if TARGET_IS_EXTENSION
-        UIApplication.shared.keyWindow?.addSubview(isPresenting ? fromView : toView)
-        #endif
+        container.window?.addSubview(isPresenting ? fromView : toView)
       }
     }