Training can fall back to CPU sometimes

[Yang-Li86]

Details as in this ChatGPT conversation: https://chatgpt.com/share/684b56a0-3778-8000-84e5-2fed4988b382

[Yang-Li86]

This happened once so far, but we need more robust program so that we make sure this does not occur in the future.

[Yang-Li86]

This is causing the aggregation to fail. There are two ways we can get around this:

1. Use the main branch (old version).
2. Try the Flwr simulation engine, maybe under the hood, it handles CUDA registration in an organized way.

For now, check job logs now and then to make sure things are training on the GPU and nothing is failing.

[Yang-Li86]

Make sure do not delete the logs so that we can load model checkpoints to resume training.