github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 07:28:15 -05:00
Code Issues Packages Projects Releases Wiki Activity

small cleanups (#9947)

Browse Source
This commit is contained in:
Francis Lata
2025-04-20 20:33:20 -04:00
committed by GitHub
parent e8910540f6
commit ea4cb2c715
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
examples/mlperf/model_train.py
View File

@@ -432,7 +432,8 @@ def train_retinanet():
model = retinanet.RetinaNet(backbone, num_classes=NUM_CLASSES)
params = get_parameters(model)
for p in params: p.to_(GPUS)
if len(GPUS) > 1:
for p in params: p.to_(GPUS)
step_times, start_epoch = [], 0
@@ -446,8 +447,7 @@ def train_retinanet():
# ** lr scheduler **
config["steps_in_train_epoch"] = steps_in_train_epoch = round_up(len(train_dataset.imgs.keys()), BS) // BS
config["steps_in_val_epoch"] = steps_in_val_epoch = (round_up(len(val_dataset.imgs.keys()), BS) // BS)
start_iter = start_epoch * steps_in_train_epoch
config["steps_in_val_epoch"] = steps_in_val_epoch = (round_up(len(val_dataset.imgs.keys()), EVAL_BS) // EVAL_BS)
# ** initialize wandb **
if (WANDB:=getenv("WANDB")):