Possédant la double nationalité hispano-sénégalaise, Keita Baldé décide d'évoluer pour le Sénégal. Il dispute son premier match avec les lions du Sénégal le contre le Niger (match de qualification pour la Coupe d'Afrique des nations de 2017 au Gabon). Ce jour-là, le Sénégal l'emporte sur le score de 2-0. Le samedi , lors de la sixième et dernière journée des qualifications pour la coupe d'Afrique des nations de 2017 au Gabon, Diao Baldé ouvre le score pour le Sénégal contre la Namibie. Après un raté de Younousse Sankharé, le joueur de la Lazio marque dans le but vide à la , inscrivant ainsi son premier but en sélection.