increase stacklevel of warnings
@@ -54,7 +54,7 @@ def warn_if_nan(value, msg=""):
if torch.is_tensor(value) and value.requires_grad:
value.register_hook(lambda x: warn_if_nan(x, msg))
if torch_isnan(value):
- warnings.warn("Encountered NaN{}".format((': ' if msg else '.') + msg))
+ warnings.warn("Encountered NaN{}".format((': ' if msg else '.') + msg), stacklevel=2)
def warn_if_inf(value, msg="", allow_posinf=False, allow_neginf=False):
@@ -66,10 +66,10 @@ def warn_if_inf(value, msg="", allow_posinf=False, allow_neginf=False):
value.register_hook(lambda x: warn_if_inf(x, msg, allow_posinf, allow_neginf))
if (not allow_posinf) and (value == float('inf') if isinstance(value, numbers.Number)
else (value == float('inf')).any()):
- warnings.warn("Encountered +inf{}".format((': ' if msg else '.') + msg))
+ warnings.warn("Encountered +inf{}".format((': ' if msg else '.') + msg), stacklevel=2)
if (not allow_neginf) and (value == -float('inf') if isinstance(value, numbers.Number)
else (value == -float('inf')).any()):
- warnings.warn("Encountered -inf{}".format((': ' if msg else '.') + msg))
+ warnings.warn("Encountered -inf{}".format((': ' if msg else '.') + msg), stacklevel=2)
def save_visualization(trace, graph_output):
