Update train.py

This commit is contained in:
Mingshuang Luo 2022-04-11 21:10:17 +08:00 committed by GitHub
parent 7320e0ea89
commit 8dcf9644cb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -959,7 +959,7 @@ def scan_pessimistic_batches_for_oom(
sp=sp, sp=sp,
batch=batch, batch=batch,
is_training=True, is_training=True,
warmup = 0.0 warmup=0.0
) )
loss.backward() loss.backward()
optimizer.step() optimizer.step()