The error points towards a device mismatch:
RuntimeError: Expected object of device type cuda but got device type cpu for argument #1 'self' in call to _th_mm

amp needs a GPU to run properly, so you would need to call .to('cuda') on the model and input.
Also, you might need to register self.n_ as a buffer via:
self.register_buffer('n_', ...)

so that it will be also pushed to the device.
Your code doesn’t provide the definition of format_frontend_output so that I cannot verify, if it’s working after these fixes.