protect against _valid_batch_iter returning an empty batch

lopuhin · Jul 28, 2019 · 27ea69e · 27ea69e
1 parent 33e561f
commit 27ea69e
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/lm/main.py b/lm/main.py
@@ -227,6 +227,8 @@ def get_valid_loss():
         with torch.no_grad():
             for ctx in _valid_batch_iter(
                     valid_dataset, batch_size=batch_size, n_ctx=n_ctx):
+                if not ctx:
+                    continue
                 ctx = torch.LongTensor(ctx).to(device)
                 logits = model(ctx)['logits']
                 loss = loss_fn(logits, ctx)