From ded167148336f8a3cc0059f6d26779255013b7f0 Mon Sep 17 00:00:00 2001
From: Dongyoon Han <dongyoon.han@navercorp.com>
Date: Mon, 7 Jun 2021 23:08:55 +0900
Subject: [PATCH] Fix stochastic depth working only with a shortcut

---
 timm/models/rexnet.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/timm/models/rexnet.py b/timm/models/rexnet.py
index 7ab8d659..462ad8fe 100644
--- a/timm/models/rexnet.py
+++ b/timm/models/rexnet.py
@@ -89,10 +89,11 @@ class LinearBottleneck(nn.Module):
             x = self.se(x)
         x = self.act_dw(x)
         x = self.conv_pwl(x)
-        if self.drop_path is not None:
-            x = self.drop_path(x)
         if self.use_shortcut:
-            x[:, 0:self.in_channels] += shortcut
+            if self.drop_path is not None:
+                x = self.drop_path(x)
+            
+	    x[:, 0:self.in_channels] += shortcut
         return x