From 0e085ac3a904e3c6623289a5bf36a8d98c1fce4c Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Tue, 18 May 2021 08:26:11 +0200
Subject: [PATCH] added unit interval domain map

---
 pts/modules/distribution_output.py | 1 +
 1 file changed, 1 insertion(+)

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()