Suppress spurious lr_scheduler.step() warning from PyTorch

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-21 18:24:54 +03:00
parent 517b87d3d8
commit fa32b2e67f

View File

@@ -11,6 +11,7 @@ import json
import logging
import math
import time
import warnings
from dataclasses import dataclass, field
from pathlib import Path
@@ -317,10 +318,13 @@ def train(cfg: TrainConfigGTAUAV) -> None:
steps_per_epoch = len(train_loader)
total_steps = cfg.epochs * steps_per_epoch
warmup_steps = cfg.warmup_epochs * steps_per_epoch
scheduler = LambdaLR(
optimizer,
lr_lambda=_cosine_warmup_schedule(warmup_steps, total_steps),
)
with warnings.catch_warnings():
warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*")
scheduler = LambdaLR(
optimizer,
lr_lambda=_cosine_warmup_schedule(warmup_steps, total_steps),
last_epoch=-1,
)
scaler = GradScaler(enabled=cfg.use_amp)
# Restore optimizer/scheduler/loss state on resume.
@@ -385,7 +389,9 @@ def train(cfg: TrainConfigGTAUAV) -> None:
)
scaler.step(optimizer)
scaler.update()
scheduler.step()
with warnings.catch_warnings():
warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*")
scheduler.step()
for key, val in loss_dict.items():
agg[key] = agg.get(key, 0.0) + float(val.item())