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
13 def exception_hook(exc_type, exc_value, tb):
14 r"""Hacks the call stack message to show all the local variables in
15 case of RuntimeError or ValueError, and prints tensors as shape,
20 repr_orig = Tensor.__repr__
21 Tensor.__repr__ = lambda x: f"{x.size()}:{x.dtype}:{x.device}"
24 print("--------------------------------------------------\n")
25 filename = tb.tb_frame.f_code.co_filename
26 name = tb.tb_frame.f_code.co_name
27 line_no = tb.tb_lineno
28 print(f' File "{filename}", line {line_no}, in {name}')
29 print(open(filename, "r").readlines()[line_no - 1])
31 if exc_type in {RuntimeError, ValueError}:
32 for n, v in tb.tb_frame.f_locals.items():
38 Tensor.__repr__ = repr_orig
40 print(f"{exc_type.__name__}: {exc_value}")
43 sys.excepthook = exception_hook
45 ######################################################################
47 if __name__ == "__main__":
57 mmm = torch.randn(2, 3)