weight_decay fix

2019-08-29 15:44:15 +02:00 · 2019-08-29 15:44:15 +02:00 · 408baf66e2
parent 7d9ffe6d4e
commit 408baf66e2
1 changed files with 3 additions and 0 deletions
--- a/train.py
+++ b/train.py
@ -258,6 +258,9 @@ def train():

            # Compute loss
            loss, loss_items = compute_loss(pred, targets, model)
+            if torch.isnan(loss):
+                print('WARNING: nan loss detected, skipping batch ', loss_items)
+                continue

            # Scale loss by nominal batch_size of 64
            loss *= batch_size / 64