Make the test use less memory.

It still use ~10G!
pep8
Make the dataset SmallNORB implement the stop parameter and use it in tests.

This lower this test memory usage from 12G to 1.6G