דיאטה היה חבר בסגל נבחרת סנגל לגביע העולם בכדורגל עד גיל 20 ב-2017 וכן בגביע האומות של אפריקה באותה שנה. את הופעת הבכורה שלו בנבחרת הבוגרת רשם במרץ 2019 במוקדמות גביע האומות של אפריקה, במשחק מול מוזמביק. בטורניר גביע האומות כבש במשחק הפתיחה בניצחון 0–2 על טנזניה. הנבחרת סיימה במקום השני ודיאטה נבחר לשחקן הצעיר המצטיין של האליפות.