[BUG] MaskedOneHotCategorical missing mode and deterministic_sample properties.

[MorganeAyle]

Describe the bug

The MaskedOneHotCategorical distribution is missing the mode and deterministic_sample properties.

Reason and Possible fixes

The MaskedOneHotCategorical distribution should overwrite the mode property to account for the conversion to and from one-hot encoding. Additionally, it should implement the deterministic_sample property. Both should be the same as those implemented in OneHotCategorical, i.e.:

@property
def mode(self) -> torch.Tensor:
    if hasattr(self, "logits"):
        return (self.logits == self.logits.max(-1, True)[0]).to(torch.long)
    else:
        return (self.probs == self.probs.max(-1, True)[0]).to(torch.long)

@property
def deterministic_sample(self):
    return self.mode

Checklist

• I have checked that there is no similar issue in the repo (required)
• I have read the documentation (required)
• I have provided a minimal working example to reproduce the bug (required)