diff --git a/examples/die.py b/examples/die.py
index 1ba5262..e51534e 100644
--- a/examples/die.py
+++ b/examples/die.py
@@ -16,7 +16,7 @@ def create_die_dtmc():
     dtmc.add_self_loops()
 
     # test if state deletion works
-    dtmc.delete_state(dtmc.get_state_by_id(1), True, True)
+    # dtmc.delete_state(dtmc.get_state_by_id(1), True, True)
 
     return dtmc
 
diff --git a/examples/monty_hall_pomdp.py b/examples/monty_hall_pomdp.py
index 251ee77..6324925 100644
--- a/examples/monty_hall_pomdp.py
+++ b/examples/monty_hall_pomdp.py
@@ -71,6 +71,7 @@ def create_monty_hall_pomdp():
     for state in pomdp.states.values():
         state.set_observation(0)
 
+    # test if the normalize function works
     pomdp.normalize()
 
     return pomdp
diff --git a/stormvogel/mapping.py b/stormvogel/mapping.py
index bc4e852..110db3a 100644
--- a/stormvogel/mapping.py
+++ b/stormvogel/mapping.py
@@ -57,7 +57,7 @@ def add_labels(model: stormvogel.model.Model) -> stormpy.storage.StateLabeling:
         Takes a model creates a state labelling object that determines which states get which labels in the stormpy representation
         """
 
-        state_labeling = stormpy.storage.StateLabeling(len(list(model.states.keys())))
+        state_labeling = stormpy.storage.StateLabeling(len(model.states))
 
         # we first add all the different labels
         for label in model.get_labels():
diff --git a/stormvogel/model.py b/stormvogel/model.py
index ce0dbd2..196f742 100644
--- a/stormvogel/model.py
+++ b/stormvogel/model.py
@@ -411,6 +411,7 @@ def is_well_defined(self) -> bool:
             ModelType.CTMC,
             ModelType.MA,
         ):
+            # TODO make it work for these models
             raise RuntimeError("Not implemented")
 
         return True
@@ -422,15 +423,19 @@ def normalize(self):
                 for action in state.available_actions():
                     sum_prob = 0
                     for tuple in state.get_outgoing_transitions(action):
-                        if isinstance(tuple[0], float) or isinstance(
-                            tuple[0], Fraction
+                        if (
+                            isinstance(tuple[0], float)
+                            or isinstance(tuple[0], Fraction)
+                            or isinstance(tuple[0], int)
                         ):
                             sum_prob += tuple[0]
 
                     new_transitions = []
                     for tuple in state.get_outgoing_transitions(action):
-                        if isinstance(tuple[0], float) or isinstance(
-                            tuple[0], Fraction
+                        if (
+                            isinstance(tuple[0], float)
+                            or isinstance(tuple[0], Fraction)
+                            or isinstance(tuple[0], int)
                         ):
                             normalized_transition = (
                                 tuple[0] / sum_prob,
@@ -444,6 +449,7 @@ def normalize(self):
             ModelType.CTMC,
             ModelType.MA,
         ):
+            # TODO make it work for these models
             raise RuntimeError("Not implemented")
 
     def __free_state_id(self):
@@ -537,6 +543,7 @@ def delete_state(
                 if state in self.markovian_states:
                     self.markovian_states.remove(state)
 
+            # we reassign the ids if specified to do so
             if reassign_ids:
                 self.states = {
                     new_id: value
@@ -562,6 +569,7 @@ def delete_state(
                         )
                     }
 
+            # we normalize the model if specified to do so
             if normalize:
                 self.normalize()