3 # Any copyright is dedicated to the Public Domain.
4 # https://creativecommons.org/publicdomain/zero/1.0/
6 # Written by Francois Fleuret <francois@fleuret.org>
8 from torch import Tensor
12 def exception_hook(exc_type, exc_value, tb):
13 r'''Hacks the call stack message in case of RuntimeError to show all
14 the local variables, and indicate for every tensor involved its
15 shape, dtype and device.
19 repr_orig=Tensor.__repr__
20 Tensor.__repr__=lambda x: f'{x.size()}:{x.dtype}:{x.device}'
23 print('--------------------------------------------------\n')
24 filename = tb.tb_frame.f_code.co_filename
25 name = tb.tb_frame.f_code.co_name
26 line_no = tb.tb_lineno
27 print(f' File "{filename}", line {line_no}, in {name}')
28 print(open(filename, 'r').readlines()[line_no-1])
30 if exc_type in { RuntimeError, ValueError }:
31 for n,v in tb.tb_frame.f_locals.items():
37 Tensor.__repr__=repr_orig
39 print(f'{exc_type.__name__}: {exc_value}')
41 sys.excepthook = exception_hook
43 ######################################################################
45 if __name__ == '__main__':