def exception_hook(exc_type, exc_value, tb):
- 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.
+ r"""Hacks the call stack message to show all the local variables
+ in case of relevant error, and prints tensors as shape, dtype and
+ device.
"""
print(f' File "{filename}", line {line_no}, in {name}')
print(open(filename, "r").readlines()[line_no - 1])
- if exc_type in {RuntimeError, ValueError}:
+ if exc_type in {RuntimeError, ValueError, IndexError}:
for n, v in tb.tb_frame.f_locals.items():
print(f" {n} -> {v}")
######################################################################
if __name__ == "__main__":
-
import torch
def dummy(a, b):