Sample from a greedy strategy with given q-value histogram

def sample(self, histogram_logits):
        """ Sample from a greedy strategy with given q-value histogram """
        histogram_probs = histogram_logits.exp()  # Batch size * actions * atoms
        atoms = self.support_atoms.view(1, 1, self.atoms)  # Need to introduce two new dimensions
        return (histogram_probs * atoms).sum(dim=-1).argmax(dim=1)