apex.amp uses monkey patching to transform the data to the appropriate dtype in the forward method.
Based on your code, it seems you are using a custom _forward method, which would be invisible for apex. Could you stick to the vanilla forward definition and check the results again?