From 221d54a6bceb7b7f2e55949f95747ff572ec453c Mon Sep 17 00:00:00 2001 From: Xiafei Qiu Date: Tue, 18 Jul 2023 19:12:07 +0800 Subject: [PATCH] update --- pytorch_blade/tests/disc/test_dynamo.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pytorch_blade/tests/disc/test_dynamo.py b/pytorch_blade/tests/disc/test_dynamo.py index b684082eb31..92c62637e1f 100644 --- 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()