vggish_classifier.py

class NoiseClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.vgg = vggish(preprocess=False, postprocess=False)
        self.vgg.embeddings = nn.Sequential(
            nn.Linear(512 * 4 * 6 * WINDOW_MULTIPLIER, 256),
            nn.ReLU(True),
            nn.Linear(256, 256),
            nn.ReLU(True),
            nn.Linear(256, 256),
            nn.ReLU(True))
        self.fc = nn.Linear(256, 1)

    def forward(self, x):
        x = self.vgg(x)
        x = self.fc(x)
        x = torch.sigmoid(x)
        x = x.squeeze()

        return x