كتاب الجنة وصفة نعيمها وأهلها
53  
The Book of Paradise, its Description, its Bounties and its Inhabitants
(17)
Chapter: The Deceased Is Shown His Place In Paradise Or The Fire; And Confirmation Of The Torment In The Grave - We Seek Refuge With Allah From That
(17)
بَاب عَرْضِ مَقْعَدِ الْمَيِّتِ مِنْ الْجَنَّةِ أَوْ النَّارِ عَلَيْهِ وَإِثْبَاتِ عَذَابِ الْقَبْرِ وَالتَّعَوُّذِ مِنْهُ
Sahih Muslim 2871 b

Al-Bara' b. 'Azib reported that this verse:

"Allah keeps those who believe steadfast with firm word in this world and the Hereafter" was revealed in connection with the torment of the grave.
Reference : Sahih Muslim 2871b
In-book reference : Book 53, Hadith 89
USC-MSA web (English) reference : Book 40, Hadith 6866
  (deprecated numbering scheme)