From 4aebb235541fc638f6e642eb43b3801206452025 Mon Sep 17 00:00:00 2001
From: lizeyan <li_zeyan@icloud.com>
Date: Thu, 15 Jul 2021 09:51:33 +0800
Subject: [PATCH] Ignore None node in .saved_tensors (It happens in some DGL
 models).

---
 torchviz/dot.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/torchviz/dot.py b/torchviz/dot.py
index a89fc8e..fb1947d 100644
--- a/torchviz/dot.py
+++ b/torchviz/dot.py
@@ -138,6 +138,8 @@ def add_nodes(fn):
         # also note that this still works for custom autograd functions
         if hasattr(fn, 'saved_tensors'):
             for t in fn.saved_tensors:
+                if t is None:
+                    continue
                 seen.add(t)
                 dot.edge(str(id(t)), str(id(fn)), dir="none")
                 dot.node(str(id(t)), get_var_name(t), fillcolor='orange')