Hi,
I’ve just tried the colab you linked to and it passed the part you mentioned (and finished fine).
Is your runtime configured to use GPU?
When the kernel dies might be an out of memory issue, did you notice that on the RAM bar on the top right corner?