def test_loss_value(self) -> None:
        """If all values are zero, the loss should be zero."""
        bs = 512
        dim = 128
        num_samples = 2

        loss_fn = WMSELoss()
        x = torch.randn(bs * num_samples, dim)

        loss = loss_fn(x)
        assert loss > 0
If all values are zero, the loss should be zero.
tests/loss/test_wmse_loss.py