def test_state_ops_cannot_removed(
        self,
        float32_buffer_not_returned_model,
        float32_buffer_not_returned_model_2,
        compute_unit,
    ):
        """
        Check the coreml_update_state should not be removed by dead_code_elimination pass.
        """
        # Test case 1
        prog = self.convert_state_model(float32_buffer_not_returned_model, "milinternal")
        assert get_op_types_in_program(prog) == [
            "identity",
            "read_state",
            "add",
            "coreml_update_state",
            "read_state",
            "add",
            "coreml_update_state",
        ]
        mlmodel = self.convert_state_model(
            float32_buffer_not_returned_model, "mlprogram", compute_unit
        )
        expected_ops = [
            "identity",
            "read_state",
            "add",
            "write_state",
            "read_state",
            "add",
            "write_state",
        ]
        # This model is failing on CPU_AND_NE, which is tracked by
        # rdar://130912134 ([Bug][Stateful model][CPU_AND_NE] Stateful model fails to run with compute_units=CPU_AND_NE)
        run_prediction = compute_unit != ct.ComputeUnit.CPU_AND_NE
        self.check_state_model(mlmodel, expected_ops, run_prediction)

        # Test case 2
        prog = self.convert_state_model(float32_buffer_not_returned_model_2, "milinternal")
        assert get_op_types_in_program(prog) == [
            "identity",
            "read_state",
            "add",
            "coreml_update_state",
            "read_state",
            "add",
            "coreml_update_state",
            "add",
            "coreml_update_state",
        ]
        mlmodel = self.convert_state_model(
            float32_buffer_not_returned_model_2, "mlprogram", compute_unit
        )
        expected_ops = [
            "identity",
            "read_state",
            "add",
            "write_state",
            "read_state",
            "read_state",
            "add",
            "write_state",
            "add",
            "write_state",
        ]
        # This model is failing on CPU_AND_NE, which is tracked by
        # rdar://130912134 ([Bug][Stateful model][CPU_AND_NE] Stateful model fails to run with compute_units=CPU_AND_NE)
        run_prediction = compute_unit != ct.ComputeUnit.CPU_AND_NE
        self.check_state_model(mlmodel, expected_ops, run_prediction)

        Check the coreml_update_state should not be removed by dead_code_elimination pass.
        
coremltools/converters/mil/frontend/torch/test/test_torch_stateful_model.py