On the Einstein relation in a heated granular gas