1 files changed, 4 insertions, 1 deletions
diff --git a/lib/automata.py b/lib/automata.py
index 5a420fe..39ac5df 100755
--- a/lib/automata.py
+++ b/lib/automata.py
@@ -956,7 +956,10 @@ class PTA:
         total_duration_mae = 0.
         total_energy = 0.
         total_energy_error = 0.
-        state = self.state[orig_state]
+        if type(orig_state) is State:
+            state = orig_state
+        else:
+            state = self.state[orig_state]
         param_dict = dict([[self.parameters[i], self.initial_param_values[i]] for i in range(len(self.parameters))])
         for function in trace:
             if isinstance(function[0], Transition):