From d2d570265e2093fbec63b81e9cab5f9c7349e739 Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Mon, 6 Jan 2020 15:18:26 +0100
Subject: [PATCH] MultivariateNormal expects named arguments

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

diff --git a/pts/modules/distribution_output.py b/pts/modules/distribution_output.py
index 453a569..515a1a6 100644
--- a/pts/modules/distribution_output.py
+++ b/pts/modules/distribution_output.py
@@ -222,6 +222,18 @@ class MultivariateNormalOutput(DistributionOutput):
 
         return loc, scale_tril
 
+    def distribution(
+        self, distr_args, scale: Optional[torch.Tensor] = None
+    ) -> Distribution:
+        loc, scale_tri = distr_args
+        distr = MultivariateNormal(loc=loc, scale_tril=scale_tri)
+
+        if scale is None:
+            return distr
+        else:
+            return TransformedDistribution(distr, [AffineTransform(loc=0, scale=scale)])
+
+
     @property
     def event_shape(self) -> Tuple:
         return (self.dim,)