Make it work with GPU on Colab

[saeub]

This PR should make it possible to train models on Colab with GPU.

• Move some tensors to the correct devices so that it should work on GPU.
• Simplify the --cuda flag (no argument needed).
• Change Pickle protocol to be compatible with Google Colab.

To use it on Colab, we used this successfully:

!git clone https://github.com/<my-github-username>/atmt
%cd atmt/
!python3 train.py \
    --data <data-dir> \
    --source-lang fr \
    --target-lang en \
    --save-dir <checkpoints-dir> \
    --cuda

Note that the preprocessed data in the data/ folder was produced with a different Pickle protocol, so you will have to re-preprocess them yourself if you want to use them on Colab.

⚠️ We haven't tested translate.py with --cuda yet, it might be necessary to move some more tensors around there. ⚠️