diff --git a/pts/modules/distribution_output.py b/pts/modules/distribution_output.py
index 019a9b5..694358d 100644
--- a/pts/modules/distribution_output.py
+++ b/pts/modules/distribution_output.py
@@ -515,6 +515,7 @@ class ImplicitQuantileOutput(IndependentDistributionOutput):
         available_domain_map_cls = {
             "Positive": nn.Softplus,
             "Real": nn.Identity,
+            "Unit": nn.Softmax,
         }
         assert (
             output_domain in available_domain_map_cls.keys()