diff --git a/Sources/Swift/MarqueeLabel.swift b/Sources/Swift/MarqueeLabel.swift
index f53b5bca..0b89aa94 100644
--- a/Sources/Swift/MarqueeLabel.swift
+++ b/Sources/Swift/MarqueeLabel.swift
@@ -1161,8 +1161,8 @@ public class MarqueeLabel: UILabel {
     }
     
     override public func animationDidStop(anim: CAAnimation, finished flag: Bool) {
-        if anim is GradientSetupAnimation {
-            if let setupAnim = maskLayer?.animationForKey("setupFade") as? CABasicAnimation, finalColors = setupAnim.toValue as? [CGColorRef] {
+        if let setupAnim = anim as? GradientSetupAnimation {
+            if let finalColors = setupAnim.toValue as? [CGColorRef] {
                 maskLayer?.colors = finalColors
             }
             // Remove regardless, since we set removeOnCompletion = false