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 is_tensor, Tensor
12 def exception_hook(exc_type, exc_value, tb):
14 repr_orig=Tensor.__repr__
15 Tensor.__repr__=lambda x: f'{x.size()}:{x.dtype}:{x.device}'
18 print('--------------------------------------------------')
19 filename = tb.tb_frame.f_code.co_filename
20 name = tb.tb_frame.f_code.co_name
21 line_no = tb.tb_lineno
22 print(f' File "{filename}", line {line_no}, in {name}')
23 print(open(filename, 'r').readlines()[line_no-1], end='')
25 if exc_type is RuntimeError:
26 for n,v in tb.tb_frame.f_locals.items():
31 Tensor.__repr__=repr_orig
33 print(f'{exc_type.__name__}: {exc_value}')
35 sys.excepthook = exception_hook
37 ######################################################################
39 if __name__ == '__main__':