Simplify axis value checks (#1501)

d1d6dbc6 · Zack Kneupper · GitHub · d49e6e45 · d1d6dbc6
Unverified Commit d1d6dbc6 authored Jul 27, 2021 by Zack Kneupper Committed by GitHub Jul 27, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

torchaudio/functional/functional.py torchaudio/functional/functional.py +2 -2

No files found.
--- a/torchaudio/functional/functional.py
+++ b/torchaudio/functional/functional.py
@@ -732,7 +732,7 @@ def mask_along_axis_iid(
        Tensor: Masked spectrograms of dimensions (batch, channel, freq, time)
    """

-    if axis != 2 and axis != 3:
+    if axis not in [2, 3]:
        raise ValueError('Only Frequency and Time masking are supported')

    device = specgrams.device
@@ -774,7 +774,7 @@ def mask_along_axis(
    Returns:
        Tensor: Masked spectrogram of dimensions (channel, freq, time)
    """
-    if axis != 1 and axis != 2:
+    if axis not in [1, 2]:
        raise ValueError('Only Frequency and Time masking are supported')

    # pack batch