Prevent kernal_normalizer to change mask dtype (#1210)

2021-07-23 21:08:58 +08:00 · 2021-07-23 21:08:58 +08:00 · 0a375614ca
parent 17fa6670eb
commit 0a375614ca
1 changed files with 1 additions and 1 deletions
--- a/mmcv/ops/carafe.py
+++ b/mmcv/ops/carafe.py
@ -268,7 +268,7 @@ class CARAFEPack(nn.Module):
        mask_channel = int(mask_c / float(self.up_kernel**2))
        mask = mask.view(n, mask_channel, -1, h, w)

-        mask = F.softmax(mask, dim=2)
+        mask = F.softmax(mask, dim=2, dtype=mask.dtype)
        mask = mask.view(n, mask_c, h, w).contiguous()

        return mask