Patch VP logSNR stability in exported SemDisDiffAE code

fcdm_diffae/vp_diffusion.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import math
 
 import torch
+import torch.nn.functional as F
 from torch import Tensor
 
 
@@ -14,8 +15,8 @@ def alpha_sigma_from_logsnr(lmb: Tensor) -> tuple[Tensor, Tensor]:
     VP constraint: alpha^2 + sigma^2 = 1.
     """
     lmb32 = lmb.to(dtype=torch.float32)
-    alpha = torch.sqrt(torch.sigmoid(lmb32))
-    sigma = torch.sqrt(torch.sigmoid(-lmb32))
+    alpha = torch.exp(0.5 * F.logsigmoid(lmb32))
+    sigma = torch.exp(0.5 * F.logsigmoid(-lmb32))
     return alpha, sigma