Numerical instability (NaN's) during training with power-spherical-vae

[BerndDoser]

The training of the power-spherical-vae crashes in epoch 378 with NaN values within the beta function:

File "/local_data/doserbd/git/Spherinator/models/../external/power_spherical/power_spherical/distributions.py", line 162, in __init__
243     MarginalTDistribution(
244   File "/local_data/doserbd/git/Spherinator/models/../external/power_spherical/power_spherical/distributions.py", line 102, in __init__
245     torch.distributions.Beta(
246   File "/local_data/doserbd/miniconda3/envs/spherinator/lib/python3.10/site-packages/torch/distributions/beta.py", line 38, in __init__
247     self._dirichlet = Dirichlet(concentration1_concentration0, validate_args=validate_args)
248   File "/local_data/doserbd/miniconda3/envs/spherinator/lib/python3.10/site-packages/torch/distributions/dirichlet.py", line 54, in __init__
249     super().__init__(batch_shape, event_shape, validate_args=validate_args)
250   File "/local_data/doserbd/miniconda3/envs/spherinator/lib/python3.10/site-packages/torch/distributions/distribution.py", line 62, in __init__
251     raise ValueError(
252 ValueError: Expected parameter concentration (Tensor of shape (64, 2)) of distribution Dirichlet(concentration: torch.Size([64, 2])) to satisfy the constraint IndependentConstraint(GreaterThan(lower_bound=0.0), 1), but found invalid values:
253 tensor([[nan, 1.],
254         [nan, 1.],
255         [nan, 1.],
256         [nan, 1.],
257         [nan, 1.],
258         [nan, 1.],
...