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