Rework PIT example train.py and data.py

[sibange]

[thequilo]

Now that the maybe_add function disappeared from the data preparation, can you add an example_to_device method to the model that only transfers to the GPU those keys from the example that are required for training? You can use pt.data.batch.example_to_device for that