Fix test failures

[carshadi]

Summary

Fix training test failures related to CPU AMP configuration and the refactored learning-rate scheduler.

Changes

• Preserve explicitly configured training and validation AMP settings on CPU and configure GradScaler using the trainer’s actual device type. Note that both of these have been supported on CPU since PyTorch 2.3+
• Update the stale scheduler test to reflect that the scheduler is local to Trainer.run().
• Verify the scheduler horizon uses the total optimizer-step count and advances once per batch.

Testing

• 66 passed
• 1 subtest passed
• 2 existing ome-zarr deprecation warnings remain.

[carshadi]

The CI tests only fail because this library does not support Python 3.10. I will fix the actions in the next PR.

[anna-grim]

Looks good!