From 9cb3f352257bfab23fc95367359373fbbc2ccbc9 Mon Sep 17 00:00:00 2001 From: Roman Pope Date: Thu, 28 May 2026 02:59:48 -0400 Subject: [PATCH] Workflow: use torch.cuda.amp.GradScaler (torch 2.2.x API) --- .riahub/workflows/train.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.riahub/workflows/train.yaml b/.riahub/workflows/train.yaml index a3c1537..af8db83 100644 --- a/.riahub/workflows/train.yaml +++ b/.riahub/workflows/train.yaml @@ -321,7 +321,7 @@ jobs: # `--device cpu` from the Train step actually takes effect. # No-op if the line already uses args.device (idempotent). if [[ -f main_finetune.py ]]; then - sed -i 's|torch\.amp\.GradScaler(device="cuda")|torch.amp.GradScaler(device=args.device, enabled=(args.device != "cpu"))|' main_finetune.py + sed -i 's|torch\.amp\.GradScaler(device="cuda")|torch.cuda.amp.GradScaler(enabled=(args.device != "cpu"))|' main_finetune.py echo "Patched main_finetune.py GradScaler for CPU/GPU device parity." fi