Support Automatic Mixed Precision (AMP)

fbcotter / pytorch_wavelets

Pytorch implementation of 2D Discrete Wavelet (DWT) and Dual Tree Complex Wavelet Transforms (DTCWT) and a DTCWT based ScatterNet

Other

943 stars 146 forks source link

Open Queuecumber opened 1 year ago

Queuecumber commented 1 year ago

This PR adds the necessary decorators to the custom autograd functions to support AMP.

DTCWT forward and backward were tested on a few examples, other functions may require testing

Fixes #45

Queuecumber commented 4 months ago

@fbcotter Could you take a look at this please?