Speed of the MaskedVectorQunantization

Thanks for releasing the nice work! I test the maskedVectorQuantization module on my task. However,  the masked version is 4x slower than the version without the masker and demasker modules. Is there any suggestions to accelerate the training? Thank you in advance!