Hi,
Yes I eventually did a regression over the distribution parameters loc and scale_tril.
So my last NN layer outputs a tensor of size [batch, 9].
However, a few things to note:

The diagonal of scale_tril must be strictly positive, so I used 1 + ELU 26:

loc, tril, diag = outputs.split(3, dim=-1)
diag = 1 + nn.functional.elu(diag)
# Each of loc, tril, diag is a tensor of size [batch, 3]


I generated scale_tril the dirty way, maybe there’s a better one:

z = torch.zeros(size=[loc.size(0)], device=device)
scale_tril = torch.stack([
    diag[:, 0], z         , z,
    tril[:, 0], diag[:, 1], z,
    tril[:, 1], tril[:, 2], diag[:, 2]
], dim=-1).view(-1, 3, 3)
# scale_tril is a tensor of size [batch, 3, 3]
dist = MultivariateNormal(loc=loc, scale_tril=scale_tril)


In order to train, instead of minimizing a criterion between a distribution sample and the label (like I thought in the last comment), it is better to minimize the log-likelihood of the label with respect to the distrubution, with the .log_prob() 9 method.

loss = - dist.log_prob(labels).mean()
loss.backward()