Update train.py (#1299)

This commit is contained in:
zr_jin 2023-10-11 10:06:00 +08:00 committed by GitHub
parent 9af144c26b
commit 0d09a44930
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -703,7 +703,7 @@ def compute_loss(
if batch_idx_train >= warm_step
else 0.1 + 0.9 * (batch_idx_train / warm_step)
)
loss = params.simple_loss_scale * simple_loss + pruned_loss_scale * pruned_loss
loss = simple_loss_scale * simple_loss + pruned_loss_scale * pruned_loss
assert loss.requires_grad == is_training