diff --git a/pyphi/tpm.py b/pyphi/tpm.py
index <HASH>..<HASH> 100644
--- a/pyphi/tpm.py
+++ b/pyphi/tpm.py
@@ -95,6 +95,9 @@ def infer_edge(tpm, a, b, contexts):
         a_off, a_on = a_in_context(context)
         return tpm[a_off][b] != tpm[a_on][b]
 
+    network_size = tpm.shape[-1]
+    all_contexts = all_states(network_size - 1)
+
     return any(a_affects_b_in_context(context) for context in contexts)
 
 
@@ -103,7 +106,7 @@ def infer_cm(tpm):
     n-dimensional form.'''
 
     network_size = tpm.shape[-1]
-    all_contexts = all_states(network_size - 1)
+    all_contexts = tuple(all_states(network_size - 1))
     cm = np.empty((network_size, network_size), dtype=int)
     for a, b in np.ndindex(cm.shape):
         cm[a][b] = infer_edge(tpm, a, b, all_contexts)
Fix optimization in `tpm.infer_cm`

`all_states` is a generator...