Proper way to mask softmax/log_softmax output

Hello, everyone!

I want to ask “How do we mask softmax output from neural network?”

In some case, like reinforcement learning, we just can do some constraint actions and we will sample the action from softmax/log_softmax output. So, we need to mask the condition which it won’t happen.

When I use mask tensor like [0. 0. 0. 1. 0. 1.] (FloatTensor) to multiply softmax output, it will sometimes comes to nan/-inf. Futhermore, it will cause runtime error: cuda runtime error (59) : device-side assert.

How do we mask the softmax/log_softmax output appropriately?

Thank you.

1 Like

Typical ways include boolean indexing with the 1/0 array, using where (in master) or just clamp’ing the infs (won’t help for NaN).

Best regards


I have same issues. Masking output of softmax will sometimes comes to nan / -inf.
I found out this is because the output of softmax is too small to float type.
The output of my NN is about -30. Sum of e^(-30) is nearly 0, which makes inf and nan.

Here’s one implementation I found to be useful.