update

alibaba · Jul 18, 2023 · 221d54a · 221d54a
1 parent 6eb2963
commit 221d54a
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/pytorch_blade/tests/disc/test_dynamo.py b/pytorch_blade/tests/disc/test_dynamo.py
@@ -26,7 +26,10 @@ def test_capture(self):
         import torch._dynamo as dynamo
         import torch_blade.dynamo
         explain_out = dynamo.explain(func1, b=torch.rand([2]))
-        self.assertEqual(explain_out.graph_count, 1)
+        if type(explain_out) is tuple:
+            self.assertEqual(len(explain_out[2]), 1)
+        else:
+            self.assertEqual(explain_out.graph_count, 1)
 
 if __name__ == '__main__':
     unittest.main()