[andrew-k681]

Support non-CUDA devices (MPS, CPU) in training and evaluation

Problem

On a machine without CUDA, YOLOX training fails with errors that point at tensor
types rather than at the missing device:

ValueError: invalid type: 'torch.mps.FloatTensor'
RuntimeError: Expected all tensors to be on the same device,
              but found at least two devices, mps:0 and cpu!

Both are symptoms of CUDA being assumed rather than selected. Neither message
mentions a device, which makes this harder to diagnose than it should be.

Root causes

1. The legacy Tensor.type() string API names a device and a dtype together.
x.type(y.type()) round-trips through a string like "torch.cuda.FloatTensor".
There is no "torch.mps.FloatTensor", so the call raises. Five call sites in
yolo_head.py, boxes.py and losses.py.

2. Call sites relied on that API for device placement. Because
.type("torch.cuda.FloatTensor") moved device and dtype, replacing it with
.to(dtype) alone is not enough: the anchor grids in YOLOXHead, built by
torch.arange on CPU, then stayed on CPU while activations moved to the
accelerator. This is the subtle half of the fix and the source of the second
error above.

3. torch.cuda.* is called unconditionally in trainer.py
(set_device, amp.GradScaler, amp.autocast) and in both evaluators
(torch.cuda.HalfTensor / FloatTensor).

Changes

• New yolox/utils/device.py with get_device_type() and get_device(),
  preferring CUDA, then MPS, then CPU. Call sites now ask for "the device"
  instead of naming a vendor.
• .type(...) replaced with .to(dtype), plus explicit device= where the old
  call was also doing placement.
• torch.cuda.set_device() guarded — only CUDA has a per-process device index;
  MPS and CPU are single-device, so indexing them is meaningless.
• torch.cuda.amp.{GradScaler,autocast} replaced with the generic torch.amp
  equivalents, parameterised by device type.
• --fp16 now raises on non-CUDA instead of silently constructing a disabled
  GradScaler and training at fp32 while the logs claim otherwise.

Testing

Fine-tuned YOLOX-Nano at 640×640 on a custom single-class COCO dataset on Apple
Silicon (M4, torch 2.14, MPS backend). Training runs, loss decreases, checkpoints
save.

Not verified on CUDA — no NVIDIA GPU was available. The changes are intended
to be behaviour-preserving there: .to(dtype) and torch.amp.autocast("cuda")
are the documented equivalents of the calls they replace, and get_device()
returns cuda:{local_rank} exactly as before when CUDA is present. That should
be confirmed before merge, and I'd appreciate a review pass from someone with a
GPU to hand.

Notes

• No behaviour change is intended for existing CUDA users; the CUDA path is
  selected first and formatted identically.
• Distributed training is untouched and still CUDA-only in practice — this
  change makes single-device non-CUDA training work, it does not claim
  multi-device MPS support.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}