github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

add jit to the training loop

Browse Source
This commit is contained in:
Francis Lata
2025-01-22 12:31:29 -08:00
parent efe64ebeaf
commit 95cdbbf237
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
examples/mlperf/model_train.py
View File

@@ -385,6 +385,7 @@ def train_retinanet():
return LambdaLR(optim, _lr_lambda)
@Tensor.train()
@TinyJit
def _train_step(model, optim, lr_scheduler, x, **kwargs):
optim.zero_grad()