batched CrossEntropyLoss in pytorch
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:]