diff --git a/OpenRA.Mods.Common/Traits/Upgrades/UpgradeOnDamage.cs b/OpenRA.Mods.Common/Traits/Upgrades/UpgradeOnDamage.cs
index c171a49af8..a66f9fd180 100644
--- a/OpenRA.Mods.Common/Traits/Upgrades/UpgradeOnDamage.cs
+++ b/OpenRA.Mods.Common/Traits/Upgrades/UpgradeOnDamage.cs
@@ -44,14 +44,11 @@ namespace OpenRA.Mods.Common.Traits
 		public UpgradeOnDamage(Actor self, UpgradeOnDamageInfo info)
 		{
 			this.info = info;
-			um = self.TraitOrDefault<UpgradeManager>();
+			um = self.Trait<UpgradeManager>();
 		}
 
 		void INotifyDamageStateChanged.DamageStateChanged(Actor self, AttackInfo e)
 		{
-			if (um == null)
-				return;
-
 			if (granted && info.GrantPermanently)
 				return;