src/transformers/time_series_utils.py/NormalOutput/domain_map
class NormalOutput:    def domain_map(cls, loc: torch.Tensor, scale: torch.Tensor):
        scale = cls.squareplus(scale)
        return loc.squeeze(-1), scale.squeeze(-1)