Need help with a certain model's loss function

sherlock.h: <p>Hey guys! I am picking up Pytorch and at the same time trying to replicate the following model:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/1/6155ed8b629cbf247987125a792d4b8c7faa5157.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/6155ed8b629cbf247987125a792d4b8c7faa5157" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/6155ed8b629cbf247987125a792d4b8c7faa5157_2_464x500.jpeg" alt="image" data-base62-sha1="dT4pGFRRPO3VIpz7SXI5CjJlrO7" width="464" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/6155ed8b629cbf247987125a792d4b8c7faa5157_2_464x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/6155ed8b629cbf247987125a792d4b8c7faa5157_2_696x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/1/6155ed8b629cbf247987125a792d4b8c7faa5157.jpeg 2x" data-dominant-color="F1F0EB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">868×935 178 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>My code for the model (modified layer for toy dataset) is:</p>
<pre data-code-wrap="python"><code class="lang-python">class DAGMM(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(2, 1),
        )
        self.decoder = nn.Sequential(
            nn.Linear(1, 2),
            nn.Tanh()
        )
        self.en = nn.Sequential(
            nn.Linear(2, 3),
            nn.Tanh(),
            nn.Dropout(),
            nn.Linear(3, 4),
            nn.Softmax(dim=1)
        )

    def forward(self, X):
        Zc = self.encoder(X)
        Xr = self.decoder(Zc)
        Re = (X-Xr).norm(dim=1, p=2).view(-1, 1)
        Z = torch.cat((Zc.detach(), Re), dim=1)
        P = self.en(Z)
        return Xr, Z, P
</code></pre>
<p>And the loss function:</p>
<pre data-code-wrap="python"><code class="lang-python">def loss_fn(X, Xr, Z, P, l1=0.1, l2=0.005):
    Phi = P.mean(dim=0)

    # Make P (k, n, 1) and Z (1, n, f) to broadcast
    Mu = (P.T.unsqueeze(2) * Z.unsqueeze(0)).sum(dim=1) / (P.sum(dim=0).view(-1, 1))

    # Same idea, create an extra dimension
    cov = (
        (
            (P.T.unsqueeze(2)) * (Z.unsqueeze(0) - Mu.unsqueeze(1))
        ).transpose(1, 2) @ (Z.unsqueeze(0) - Mu.unsqueeze(1))
    ) / (P.sum(dim=0).view(-1, 1).unsqueeze(1))

    cov_inv = cov.inverse() #torch.cholesky_inverse(torch.linalg.cholesky(cov))

    Zc = Z.unsqueeze(0) - Mu.unsqueeze(1)
    E = -(
        (
            (
                (Zc.unsqueeze(2) @ cov_inv.unsqueeze(1)
            ) @ Zc.unsqueeze(-1)).squeeze(-1).exp()
        ) * (Phi / (np.pi * cov.det()).sqrt()).view(-1, 1).unsqueeze(1)
    ).sum(dim=0).log()

    P_cov = ((1 / cov.diagonal(dim1=1, dim2=2)).ravel() ** 2).sum().item()

    loss = (X-Xr).norm(dim=1, p=2).mean()
    loss += l1 * E.mean().item()
    loss += l2 * P_cov

    return loss
</code></pre>
<p>When I run the model, it doesn’t really improve. Given that this is the first time that I am using the framework and making vectorization attempts, I believe that it is probably due to my loss function code.</p>
<p>If you also notice other problems or a best practice that I could adopt, please do tell!</p>