Invalid gradient at index 0 -- while using an custom layer with amp

Hi, I am trying to train my network with Apex to get any speed up using FP16 (opt_level=O1). The network contains a custom layer written using C++/CUDA available from here:
https://github.com/CharlesShang/DCNv2/blob/master/dcn_v2.py

and specifically being used in the following architecture:
https://github.com/xinntao/EDVR/tree/master/codes

Following the suggestions in #345, I have registered the forward and backward functions in the above to be specifically of FP32 precision using '@amp.float_function'. However, when I try to train the network, I get the following error at the 'scaled_loss.backward()' computation:

"RuntimeError: Function _DCNv2Backward returned an invalid gradient at index 0 - expected type torch.cuda.HalfTensor but got torch.cuda.FloatTensor"

As per my understanding, by registering the function specifically with FP32, we are asking Apex to leave that layer alone in FP32 and once it returns the gradient in FP32 the 'scaled_loss.backward()' should be able to take it as it is, since, it is supposed to be in FP32 anyway in opt_level=O1. Why is it complaining for tensor being a Float?

Any advise in this regard would be very helpful.
Thanks!

Start with the custom layer in dcn_v2.py and the EDVR architecture, then trace the call to scaled_loss.backward() and the use of amp.float_function described in #345. Reproduce the invalid _DCNv2Backward gradient error and determine the expected gradient behavior for this FP32-registered layer; done means the error is explained and a verified resolution or documentation is available.