How to fix the error "cannot import name 'custom_fwd' from 'torch.amp‘" in ska.py

[zhqshen]

The system report error "cannot import name 'custom_fwd' from 'torch.amp‘" due to the call "from torch.amp import custom_fwd, custom_bwd" in ska.py, so,how to fix thie error?