distributions.Categorical fails with constraint.Simplex but manual check passes

Hey @KFrank,

thanks for the quick response. I dumped the tensor in a try except block. I uploaded it here:
https://oc.embl.de/index.php/s/IxhtqcoNbUnpLVS

and the behaviour on my machine as

The relevant parts of my environment are:

Name Version Build Channel

cudatoolkit 11.3.1 h9edb442_10 conda-forge
python 3.10.5 h582c2e5_0_cpython conda-forge
python_abi 3.10 2_cp310 conda-forge
pytorch 1.12.0 py3.10_cuda11.3_cudnn8.3.2_0 pytorch
pytorch-lightning 1.6.5 pypi_0 pypi
pytorch-mutex 1.0 cuda pytorch

complete env here: https://oc.embl.de/index.php/s/UwFFBFDJy6s71t0