In this work, we propose an attention-based adaptive optics method that uses a non-local block to integrate phase diversity with a convolutional neural network (CNN). The simulation results showcase the effectiveness of the proposed method to mitigate the ambiguity problem of phase retrieval and better performance than traditional CNN-based wavefront correction.