diff --git a/deckard/base/attack/attack.py b/deckard/base/attack/attack.py
index 5a02e0ce..45d8e9b0 100644
--- a/deckard/base/attack/attack.py
+++ b/deckard/base/attack/attack.py
@@ -213,11 +213,13 @@ def __call__(
                 end = process_time_ns()
             try:
                 start = process_time_ns()
+                start_timestamp = time()
                 adv_probabilities = model.predict_proba(samples)
                 end = process_time_ns()
-
+                end_timestamp = time()
             except AttributeError:
                 start = process_time_ns()
+                start_timestamp = time()
                 adv_probabilities = model.predict(samples)
                 end = process_time_ns()
                 end_timestamp = time()
diff --git a/deckard/base/model/model.py b/deckard/base/model/model.py
index 83722b94..3373d0a0 100644
--- a/deckard/base/model/model.py
+++ b/deckard/base/model/model.py
@@ -561,6 +561,8 @@ def predict(self, data=None, model=None, predictions_file=None):
             start = process_time_ns()
             start_timestamp = time()
             predictions = model.predict(data[1])
+            end = process_time_ns() - start
+            end_timestamp = time()
         except NotFittedError as e:  # pragma: no cover
             logger.warning(e)
             logger.warning(f"Model {model} is not fitted. Fitting now.")