I assume you are running out of memory, if you try to pass the whole dataset into the model?
If so, have a look at the Data loading tutorial to see, how a Dataset and DataLoader are working and how batches are created.
The model itself is quite small and its parameters will just take approx. 1.5MB of memory.