Torch `autocast` deprecation warning

[benjaminking]

(Copied from a comment in #670 and made into its own issue to allow for an agent assignment)

Torch warning at the start of training:

1740758286198 aqua-gpu-dallas:gpu2 DEBUG [INFO|trainer.py:2314] 2025-02-28 15:58:01,402 >> ***** Running training *****
[INFO|trainer.py:2315] 2025-02-28 15:58:01,402 >>   Num examples = 9,365
[INFO|trainer.py:2316] 2025-02-28 15:58:01,402 >>   Num Epochs = 35
[INFO|trainer.py:2317] 2025-02-28 15:58:01,402 >>   Instantaneous batch size per device = 64
[INFO|trainer.py:2319] 2025-02-28 15:58:01,402 >>   Training with DataParallel so batch size has been adjusted to: 32
[INFO|trainer.py:2320] 2025-02-28 15:58:01,402 >>   Total train batch size (w. parallel, distributed & accumulation) = 64
[INFO|trainer.py:2321] 2025-02-28 15:58:01,402 >>   Gradient Accumulation steps = 2
[INFO|trainer.py:2322] 2025-02-28 15:58:01,402 >>   Total optimization steps = 5,000
[INFO|trainer.py:2323] 2025-02-28 15:58:01,404 >>   Number of trainable parameters = 1,370,638,336

  0% 0/5000 [00:01<?, ?it/s]
/usr/local/lib/python3.10/dist-packages/torch/utils/checkpoint.py:295: FutureWarning:

`torch.cpu.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cpu', args...)` instead.

[benjaminking]

We will plan to tie this upgrade to the upgrade of the Python and Ubuntu version.