diff --git a/MisterFusion.podspec b/MisterFusion.podspec
index 2296d18..10f6d4f 100644
--- a/MisterFusion.podspec
+++ b/MisterFusion.podspec
@@ -9,7 +9,7 @@
 Pod::Spec.new do |s|
   s.name             = "MisterFusion"
 
-  s.version          = "4.0.0"
+  s.version          = "4.0.1"
 
   s.summary          = "MisterFusion makes more easier to use AutoLayout in Swift & Objective-C code."
 
diff --git a/MisterFusion/MisterFusion.swift b/MisterFusion/MisterFusion.swift
index 3d476d1..50016b0 100644
--- a/MisterFusion/MisterFusion.swift
+++ b/MisterFusion/MisterFusion.swift
@@ -342,7 +342,7 @@ extension View {
         let item: _LayoutObject = misterFusion.item ?? self
 
         #if os(iOS) || os(tvOS)
-        let traitCollection = UIApplication.shared.keyWindow?.traitCollection
+        let traitCollection = window?.traitCollection
         if let horizontalSizeClass = misterFusion.horizontalSizeClass
             , horizontalSizeClass != traitCollection?.horizontalSizeClass {
             return nil