Hi, I have a custom torch.autograd.Function function and try to use mix precision training but I got error  below. I use pytorch 1.7
NameError: name ‘custom_fwd’ is not defined
Here is the example code
class MyFloat32Func(torch.autograd.Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, input):
        ctx.save_for_backward(input)
        pass
        return fwd_output
    @staticmethod
    @custom_bwd
    def backward(ctx, grad):
        pass
    
func=MyFloat32Func.apply
func(2)

And I got error
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-3-744176e249fa> in <module>
----> 1 class MyFloat32Func(torch.autograd.Function):
      2     @staticmethod
      3     @custom_fwd(cast_inputs=torch.float32)
      4     def forward(ctx, input):
      5         ctx.save_for_backward(input)

<ipython-input-3-744176e249fa> in MyFloat32Func()
      1 class MyFloat32Func(torch.autograd.Function):
      2     @staticmethod
----> 3     @custom_fwd(cast_inputs=torch.float32)
      4     def forward(ctx, input):
      5         ctx.save_for_backward(input)

NameError: name 'custom_fwd' is not defined

If I call the following print
print(torch.cuda.amp.custom_fwd)
print(torch.cuda.amp.custom_fwd(cast_inputs=torch.float32))

I get
<function custom_fwd at 0x7f9b0310cef0>
functools.partial(<function custom_fwd at 0x7f9b0310cef0>, cast_inputs=torch.float32)