Hello, I am new to Pyro and probabilistic programming. I just finished reading the VAE example in the docs and feel confident in writing probabilistic models in Pyro. However, I decided to implement one using PyTorch distributions for the sake of learning. I want to know if this model is right or wrong.
import torch as T import torch.nn as nn import torch.distributions as dist class NN(nn.Module): def __init__(self): super().__init__() self.lin = nn.Linear(768, 128) self.fc = nn.Linear(128, 64) def forward(self, x): x = self.lin(x) loc = self.fc(x) scale = nn.functional.softplus(self.fc(x)) return loc, scale def model(): # implementing p(z|x) x = T.rand([1, 1, 28, 28]) net = NN() loc, scale = net(x) return dist.Normal(loc, scale).sample()