more changes

1 files changed, 2 insertions, 1 deletions

diff --git a/src/cascade_creation.py b/src/cascade_creation.py
index 93ce662..39fc531 100644
--- a/src/cascade_creation.py
+++ b/src/cascade_creation.py
@@ -1,6 +1,7 @@
 import networkx as nx
 import numpy as np
 import collections
+import timeout
 #from itertools import izip
 from sklearn.preprocessing import normalize
 
@@ -70,6 +71,7 @@ class Cascade(list):
         return candidate_infectors
 
 
+@timeout.timeout(5)
 def icc_cascade(G, p_init):
     """
     Returns boolean vectors for one cascade
@@ -87,7 +89,6 @@ def icc_cascade(G, p_init):
         active = active & susceptible
         susceptible = susceptible & np.logical_not(active)
     if not cascade:
-        print("Empty cascade, consider changing p_init or n_nodes. Retrying.")
         return icc_cascade(G, p_init)
     return cascade