for word, i in word_index.items():
        try:
            embedding_matrix[i] = embedding_index[word]
        except KeyError:
            try:
                embedding_matrix[i] = embedding_index[word.lower()]
            except KeyError:
                try:
                    embedding_matrix[i] = embedding_index[word.title()]
                except KeyError:
                    unknown_words.append(word)

    return embedding_matrix, unknown_words

class SpatialDropout(nn.Dropout2d):
    def forward(self, x):
        x = x.unsqueeze(2)  # (N, T, 1, K)
        x = x.permute(0, 3, 2, 1)  # (N, K, 1, T)
        x = super(SpatialDropout, self).forward(
            x
        )  # (N, K, 1, T), some features are masked
        x = x.permute(0, 3, 2, 1)  # (N, T, 1, K)
        x = x.squeeze(2)  # (N, T, K)
        return x