Warn when the ever-present PyTorch assertion error occurs on backward()

JohnVinyard · Mar 3, 2018 · ade66d8 · ade66d8
1 parent ffd7c25
commit ade66d8
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/zounds/learn/supervised.py b/zounds/learn/supervised.py
@@ -1,5 +1,6 @@
 from trainer import Trainer
 import numpy as np
+import warnings
 
 
 class SupervisedTrainer(Trainer):
@@ -96,8 +97,12 @@ def batch(d, l, test=False):
                 minibatch_data = data[minibatch_slice]
                 minibatch_labels = labels[minibatch_slice]
 
-                inp, output, e = batch(
-                    minibatch_data, minibatch_labels, test=False)
+                try:
+                    inp, output, e = batch(
+                        minibatch_data, minibatch_labels, test=False)
+                except RuntimeError as e:
+                    warnings.warn(e.message)
+                    continue
 
                 # test batch
                 if test_size: