iteratively plot numpy or torch arrays as images and switch on keypress

Question

i came across this post that shows how to iteratively plot random plots by kyepress : link here

it works fine.

i can\'t modify it so my tensor of shape (64,64,