code: def sample(self, histogram_logits):
          
          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)
docs: Sample from a greedy strategy with given q-value histogram