Thank you very much! @torch.cuda.amp.custom_fwd works!!
class MyFloat32Func(torch.autograd.Function):
    @staticmethod
    @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return torch.tensor(input,dtype=torch.float32)
    @staticmethod
    @torch.cuda.amp.custom_fwd
    def backward(ctx, grad):
        return 1
    
func=MyFloat32Func.apply
func(2)

Just wondering, why import custom_fwd from torch.cuda.amp does not work…
from torch.cuda.amp import custom_fwd


class MyFloat32Func(torch.autograd.Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return torch.tensor(input,dtype=torch.float32)
    @staticmethod
    @custom_bwd
    def backward(ctx, grad):
        return 1
    
func=MyFloat32Func.apply
func(2)

Here is the error
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-12-d1461ca0ff6e> in <module>
      2 
      3 
----> 4 class MyFloat32Func(torch.autograd.Function):
      5     @staticmethod
      6     @custom_fwd(cast_inputs=torch.float32)

<ipython-input-12-d1461ca0ff6e> in MyFloat32Func()
      9         return torch.tensor(input,dtype=torch.float32)
     10     @staticmethod
---> 11     @custom_bwd
     12     def backward(ctx, grad):
     13         return 1

NameError: name 'custom_bwd' is not defined