From 65f86ba47555a07a0fd6483633c5a502e77f4b18 Mon Sep 17 00:00:00 2001 From: Angela Yi Date: Fri, 29 Mar 2024 16:34:15 -0700 Subject: [PATCH] Update tests to also count constants Differential Revision: D55545960 --- exir/backend/test/test_partitioner.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/exir/backend/test/test_partitioner.py b/exir/backend/test/test_partitioner.py index c7cc2b300a3..828518fa96d 100644 --- a/exir/backend/test/test_partitioner.py +++ b/exir/backend/test/test_partitioner.py @@ -232,7 +232,9 @@ def partition( # Check Owning Program still owns all constant data owning_program = delegated.exported_program() - self.assertEqual(len(owning_program.state_dict), 3) + self.assertEqual( + len(owning_program.state_dict) + len(owning_program.constants), 3 + ) self.assertEqual(len(owning_program.graph_signature.buffers), 2) self.assertEqual(len(owning_program.graph_signature.parameters), 1) @@ -321,7 +323,7 @@ def partition( delegated.exported_program().graph_module, lowered_module_node.name ) delegated_ep = lower_module.original_module - self.assertEqual(len(delegated_ep.state_dict), 3) + self.assertEqual(len(delegated_ep.state_dict) + len(delegated_ep.constants), 3) self.assertEqual(len(delegated_ep.graph_signature.buffers), 2) self.assertEqual(len(delegated_ep.graph_signature.parameters), 1) @@ -375,7 +377,9 @@ def partition( # Check Owning Program still owns only buffers owning_program = delegated.exported_program() - self.assertEqual(len(owning_program.state_dict), 2) + self.assertEqual( + len(owning_program.state_dict) + len(owning_program.constants), 2 + ) self.assertEqual(len(owning_program.graph_signature.buffers), 2) self.assertEqual(len(owning_program.graph_signature.parameters), 0)