OK, Thank God! I finally got the hang of it! here is what I ended up doing !
# display a 2D manifold of the digits
embeddingsize = model.embedding_size
# figure with 20x20 digits
n = 20  
digit_size = 28

z1 = torch.linspace(-2, 2, n)
z2 = torch.linspace(-2, 2, n)

z_grid = np.dstack(np.meshgrid(z1, z2))
z_grid = torch.from_numpy(z_grid).to(device)
z_grid = z_grid.reshape(-1, embeddingsize)

x_pred_grid = model.decoder(z_grid)
x_pred_grid= x_pred_grid.cpu().detach().view(-1, 1, 28,28)
img = make_grid(x_pred_grid,nrow=n).numpy().transpose(1,2,0)

plt.figure(figsize=(10, 10))
plt.imshow(img)
plt.show()

and the output is :

visualization_vae_pytorch.jpg789×771 149 KB