add bn to ca

bubbliiiing · Apr 18, 2022 · 4127cd0 · 4127cd0
1 parent 926072a
commit 4127cd0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nets/attention.py b/nets/attention.py
@@ -102,7 +102,7 @@ def forward(self, x):
         x_h = torch.mean(x, dim = 3, keepdim = True).permute(0, 1, 3, 2)
         x_w = torch.mean(x, dim = 2, keepdim = True)
 
-        x_cat_conv_relu = self.relu(self.conv_1x1(torch.cat((x_h, x_w), 3)))
+        x_cat_conv_relu = self.relu(self.bn(self.conv_1x1(torch.cat((x_h, x_w), 3))))
 
         x_cat_conv_split_h, x_cat_conv_split_w = x_cat_conv_relu.split([h, w], 3)