X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=tensorstack.py;h=c9a6c2f301ad50b8ea75ff05b01dcf89b8578db9;hb=ca897077ed89fbc3c7e8d812ad262146a0c72b71;hp=8360f9b9102bff9ba1642a2f458d25130f051db1;hpb=88302d2fbe5ca4adc72f24a78e8d2abb5326418c;p=pytorch.git

diff --git a/tensorstack.py b/tensorstack.py
index 8360f9b..c9a6c2f 100755
--- a/tensorstack.py
+++ b/tensorstack.py
@@ -10,9 +10,9 @@ from torch import Tensor
 import sys
 
 def exception_hook(exc_type, exc_value, tb):
-    r'''Hacks the call stack message in case of RuntimeError to show all
-    the local variables, and indicate for every tensor involved its
-    shape, dtype and device.
+    r'''Hacks the call stack message to show all the local variables in
+    case of RuntimeError or ValueError, and prints tensors as shape,
+    dtype and device.
 
     '''
 
@@ -27,7 +27,7 @@ def exception_hook(exc_type, exc_value, tb):
         print(f'  File "{filename}", line {line_no}, in {name}')
         print(open(filename, 'r').readlines()[line_no-1])
 
-        if exc_type is RuntimeError:
+        if exc_type in { RuntimeError, ValueError }:
             for n,v in tb.tb_frame.f_locals.items():
                 print(f'  {n} -> {v}')