transforms.py

import torch

from torchaudio.transforms import Resample
from torchvision.transforms import RandomCrop



class Transform(torch.nn.Module):
    def __init__(
            self,
            input_sr,
            sample_rate,
            max_seconds,
            *args,
            **kwargs
    ):
        super().__init__()
        self.resample = Resample(orig_freq=input_sr, new_freq=sample_rate)
        self.random_crop = RandomCrop((1, int(max_seconds * sample_rate)), pad_if_needed=True)

    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
        resampled = self.resample(waveform)
        croped = self.random_crop(resampled)
        return croped