Writing Snake in 12 Lines of PyTorch

Today we are programming a version of Snake where the snake can loop around the screen. However, you can change 2 lines to produce the original version of Snake, but I will leave that as an exercise for the reader. We will be using PyTorch and NumPy. This could have been done completely in either, but I prefer the PyTorch tensor API and NumPy has a nice function called unravel_index that we will be using. I will also not be counting the imports and function declaration; call it freedom of artistic expression.