def forward(
        self,
        logits_t: torch.Tensor,
        logits_tp1: torch.Tensor,
        atoms_target_t: torch.Tensor,
    ) -> torch.Tensor:
        """Compute the loss

        Args:
            logits_t (torch.Tensor): predicted atoms at step T.
                shape: [bs; num_atoms]
            logits_tp1 (torch.Tensor): predicted atoms at step T+1.
                shape: [bs; num_atoms]
            atoms_target_t (torch.Tensor): target atoms at step T.
                shape: [bs; num_atoms]

        Returns:
            torch.Tensor: computed loss
        """
        probs_tp1 = torch.softmax(logits_tp1, dim=-1)
        tz = torch.clamp(atoms_target_t, self.v_min, self.v_max)
        tz_z = torch.abs(tz[:, None, :] - self.z[None, :, None])
        tz_z = torch.clamp(1.0 - (tz_z / self.delta_z), 0.0, 1.0)
        probs_target_t = torch.einsum("bij,bj->bi", (tz_z, probs_tp1)).detach()
        loss = _ce_with_logits(logits_t, probs_target_t).mean()
        return loss
Compute the loss

        Args:
            logits_t (torch.Tensor): predicted atoms at step T.
                shape: [bs; num_atoms]
            logits_tp1 (torch.Tensor): predicted atoms at step T+1.
                shape: [bs; num_atoms]
            atoms_target_t (torch.Tensor): target atoms at step T.
                shape: [bs; num_atoms]

        Returns:
            torch.Tensor: computed loss
        
catalyst/contrib/losses/regression.py