BlueBrain / morphoclass

Neuronal morphology preparation and classification using Machine Learning.
https://morphoclass.readthedocs.io
Apache License 2.0
8 stars 4 forks source link

Checkpoints Save/Load from Different Devices #30

Open FrancescoCasalegno opened 2 years ago

FrancescoCasalegno commented 2 years ago

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