batched CrossEntropyLoss in pytorch 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.
I solved this by using torch.where to give the irrelevent entries a really big negative value, so that they vanish after an exp. I also took advantage of the log of CEL trick.

loss = torch.zeros(1).type(dtype)

states = torch.cat([states_[_] for _ in idc[:self.batch_size]]).type(dtype)
scores = torch.tensor([scores_[_] for _ in idc[:self.batch_size]]).type(dtype).view(-1, 1)
policies = torch.cat([policies_[_] for _ in idc[:self.batch_size]]).type(dtype)

values, logits = self.network(states)
value_loss = mse(values, scores)

p = torch.nn.functional.normalize(policies, dim=[1,2,3], p=1)
logit_exp = torch.exp(torch.where(p &gt; 0, logits, inf))
s = torch.log(torch.sum(logit_exp, dim=[1,2,3])).view(-1, 1, 1, 1)
log_q = logits - torch.mul(ones, s)
policy_loss = torch.sum(torch.mul(p, -log_q))/self.batch_size

loss = value_loss + policy_loss
loss.backward()
optimizer.step()
optimizer.zero_grad()
idc = idc[self.batch_size:]