Fix index for get_base_samples in interleaved MTMVN
This was overlooked in the previous fix
@@ -76,7 +76,7 @@ class MultitaskMultivariateNormal(MultivariateNormal):
base_samples = super().get_base_samples(sample_shape)
if not self._interleaved:
# flip shape of last two dimensions
- new_shape = sample_shape + self._output_shape[:-3] + self._output_shape[:-3:-1]
+ new_shape = sample_shape + self._output_shape[:-2] + self._output_shape[:-3:-1]
return base_samples.view(new_shape).transpose(-1, -2).contiguous()
return base_samples.view(*sample_shape, *self._output_shape)
