consider following code will this part of computation graph correctly and gradients are computed correctly.I mean can distributions be part of graph

import torch
import torch.distributions.normal as tdn
dist = tdn.Normal(mean,std) #assume network predicts mean,std
loss = -1 * dist.log_prob(y)