I'm wondering how to implement this with pytorch built-ins. I've got a 3 dimensional input of uints called policy. Most of the entries are zero, and if I were to L1 normalize this I would have a (target) probability distribution.
I've also got the output of a linear layer, called 'logit', with the same shape as 'policy'. I must turn this into a probability distribution by taking the softmax, but only over the entries where policy is non-zero.
The loss is then -sum(log(logit_masked_softmax) * policy_normalized))
I've implemented this manually with the nn.functional module using boolean indexing. The problem is that I want to do this in batches, where a 4 dimensional tensor represents the batches of 3 dimensional inputs. I am convinced that there must be a built-in way to achieve this and it probably is also faster and more numerically stable.