Update tests to also count constants (#2775)

Summary: Pull Request resolved: #2775

Reviewed By: cccclai

Differential Revision: D55545960

fbshipit-source-id: 047890a443925bf473316c74d5be0a5b5867a98a

Author: angelayi

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)