Checkpoints Save/Load from Different Devices

We are now fixing the `torch.load()` by specifying the `device`. However, running `dvc repro` on a CPU vs. GPU system would produce different artifacts. Should we also modify the code of `morphoclass training` to make sure that it's consistently written as CPU?
Basically I am talking about this:
https://github.com/BlueBrain/morphoclass/blob/8fca841b715af3057db4cbfde19773674e8668da/src/morphoclass/training/training_log.py#L177-L179
But I guess it's not trivial (maybe we need to iterate through each element of the `dict` and convert it?) because there's no equivalent of `torch.load`'s `map_location=device` for `torch.save`.

_Originally posted by @FrancescoCasalegno in https://github.com/BlueBrain/morphoclass/issues/15#issuecomment-1103826655_