fix softmax cpu

unicore/modules/softmax_dropout.py

@@ -94,6 +94,7 @@ def softmax_dropout(input, dropout_prob, is_training=True, mask=None, bias=None)
         torch.Tensor: the result after softmax
     """
     input = input.contiguous()
+    if input.is_cuda and input.shape[-1] <= 2048:
         input_size = input.size()
         if mask is not None:
             _check_mask(mask, input)
@@ -102,15 +103,12 @@ def softmax_dropout(input, dropout_prob, is_training=True, mask=None, bias=None)
             _check_bias(bias, input)
             bias = bias.contiguous().view(-1, input_size[-2], input_size[-1])
         input = input.view(-1, input_size[-2], input_size[-1])
-    if input.is_cuda and input.shape[-1] <= 2048:
         return SoftmaxDropoutFast.apply(
             is_training, input, mask, bias, dropout_prob
         ).view(*input_size)
     else:
-        if mask is None:
+        if mask is not None:
             input += mask
         if bias is not None:
             input += bias
-        return F.dropout(
-            F.softmax(input, dim=-1), p=dropout_prob, training=is_training
-        ).view(*input_size)
+        return F.dropout(F.softmax(input, dim=-1), p=dropout_prob, training=is_training)