multiline KeyError msg python bug workaround (#18557)

Summary:
make multiline KeyError msg readable by working around a python bug https://bugs.python.org/issue2651

discussion: https://github.com/pytorch/pytorch/issues/16647
Pull Request resolved: https://github.com/pytorch/pytorch/pull/18557

Differential Revision: D14681086

Pulled By: soumith

fbshipit-source-id: acbd13a823302c854c3d364028ed414fd8ce6bc8

diff --git a/torch/utils/data/dataloader.py b/torch/utils/data/dataloader.py
index 97b9427..74a0c26 100644 (file)
--- a/torch/utils/data/dataloader.py
+++ b/torch/utils/data/dataloader.py
@@ -569,7 +569,12 @@ class _DataLoaderIter(object):
         self.rcvd_idx += 1
         self._put_indices()
         if isinstance(batch, _utils.ExceptionWrapper):
-            raise batch.exc_type(batch.exc_msg)
+            # make multiline KeyError msg readable by working around
+            # a python bug https://bugs.python.org/issue2651
+            if batch.exc_type == KeyError and "\n" in batch.exc_msg:
+                raise Exception("KeyError:" + batch.exc_msg)
+            else:
+                raise batch.exc_type(batch.exc_msg)
         return batch
 
     def __getstate__(self):