Sahih Muslim Book 40, Hadith Number 6866.

Chapter : The dead would be shown his seat in Paradise and hell, and the affirmation of the torment of the grave and seeking refuge from it.

Al-Bara’ b. ‘Azib reported that this verse: “Allah keeps steadfast those who believe with firm word in this world and the Hereafter,” was revealed in connection with the torment of the grave.