How to use wavelet transform in "Denoise Speech Using Deep Learning Networks" example?

I have implemented the example here, but I also want to do this example with wavelet. Or can I do it with another fft like? I would be very happy if you tell me how to integrate the wavelet.

https://www.mathworks.com/help/deeplearning/ug/denoise-speech-using-deep-learning-networks.html#mw_rtc_DenoiseSpeechUsingDeepLearningNetworksExample_9FBC1D4D