The doc of `AugMix()` explains `chain_depth` parameter wrongly

[hyperkai]

📚 The doc issue

The doc of AugMix() says about chain_depth parameter as shown below:

Parameters:
...

• chain_depth (int, optional) – The depth of augmentation chains. A negative value denotes stochastic depth sampled from the interval [1, 3]. Default is -1.

But the source code says self.chain_depth if self.chain_depth > 0 as shown below:

depth = self.chain_depth if self.chain_depth > 0 else int(torch.randint(low=1, high=4, size=(1,)).item())

Suggest a potential alternative/fix

So, the doc should say as shown below:

Parameters:
...

• chain_depth (int, optional) – The depth of augmentation chains. A zero or negative value denotes stochastic depth sampled from the interval [1, 3]. Default is 0.

Or, self.chain_depth > 0 should be changed to self.chain_depth >= 0 as shown below:

depth = self.chain_depth if self.chain_depth >= 0 else int(torch.randint(low=1, high=4, size=(1,)).item())