Suppose that you have token logits from a Transformer model. The logits are of shape (batch_size, sequence_length). In the example below, (4,512).
I would like to turn every token logit into an independent Bernouilli distribution (as I’m implementing a research paper). Next, the loss is defined as the log probability of the ground-truth labels.
import torch
logits = torch.randn((4,512))
dist_per_token = torch.distributions.Bernoulli(logits=logits) # every token is now a Bernouilli distribution
labels = torch.randint(0, 1, (4,512))
loss_per_token = dist_per_token.log_prob(labels)
However, when my labels is a LongTensor, i get the following error:
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
<ipython-input-5-196b06f84a85> in <module>()
1 labels = torch.randint(0, 1, (4,512))
2
----> 3 dist_per_token.log_prob(labels)
1 frames
/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in binary_cross_entropy_with_logits(input, target, weight, size_average, reduce, reduction, pos_weight)
2824 raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))
2825
-> 2826 return torch.binary_cross_entropy_with_logits(input, target, weight, pos_weight, reduction_enum)
2827
2828
RuntimeError: result type Float can't be cast to the desired output type Long
Shouldn’t this work, as the Bernouilli is a discrete distribution? Or should one always provide FloatTensors to log_prob? It only works when I type dist_per_token.log_prob(labels.float()).