diff --git a/bnlearn/bnlearn.py b/bnlearn/bnlearn.py
index ed72834..e03c38f 100644
--- a/bnlearn/bnlearn.py
+++ b/bnlearn/bnlearn.py
@@ -945,11 +945,13 @@ def get_edge_properties(model, color='#000000', weight=1, minscale=1, maxscale=1
     # if adjmat is not None and ('lingam' not in model['config']['method']):
     if adjmat is not None:
         if verbose>=3: print('[bnlearn]> Set edge weights based on the [%s] test statistic.' %(model['independence_test'].columns[-2]))
-        logp = -np.log10(model['independence_test']['p_value'])
-        Iloc = np.isinf(logp)
-        max_logp = np.max(logp[~Iloc]) * 1.5  # For visualization purposes, set the max higher then what is present to mark the difference.
-        if np.isnan(max_logp): max_logp = 1
-        logp.loc[Iloc] = max_logp
+        # Compute logp
+        logp = compute_logp(model['independence_test']['p_value'])
+        # logp = -np.log10(model['independence_test']['p_value'])
+        # Iloc = np.isinf(logp)
+        # max_logp = np.max(logp[~Iloc]) * 1.5  # For visualization purposes, set the max higher then what is present to mark the difference.
+        # if np.isnan(max_logp): max_logp = 1
+        # logp.loc[Iloc] = max_logp
         # Rescale the weights
         weights = _normalize_weights(logp.values, minscale=minscale, maxscale=maxscale)
         # Add to adjmat
@@ -987,6 +989,15 @@ def get_edge_properties(model, color='#000000', weight=1, minscale=1, maxscale=1
     return edges
 
 
+def compute_logp(p_value):
+    logp = -np.log10(p_value)
+    Iloc = np.isinf(logp)
+    # For visualization purposes, set the max higher then what is present to mark the difference.
+    max_logp = np.max(logp[~Iloc]) * 1.5
+    if np.isnan(max_logp): max_logp = 1
+    logp.loc[Iloc] = max_logp
+    return logp
+
 # %% PLOT
 def plot_graphviz(model,
                   edge_labels='weight',
@@ -1092,19 +1103,24 @@ def plot_graphviz(model,
         if verbose>=3: print('[bnlearn] >Setting edge labels to pvalue.')
         source = model.get('independence_test')['source']
         target = model.get('independence_test')['target']
-        logp = -np.log10(model.get('independence_test')['p_value'])
-        Iloc = np.isinf(logp)
+        # Compute logp
+        logp = compute_logp(model['independence_test']['p_value'])
+        # logp = -np.log10(model.get('independence_test')['p_value'])
+        # Iloc = np.isinf(logp)
         # For visualization purposes, set the max higher then what is present to mark the difference.
-        max_logp = np.max(logp[~Iloc]) * 1.5
-        if np.isnan(max_logp): max_logp = 1
-        logp.loc[Iloc] = max_logp
+        # max_logp = np.max(logp[~Iloc]) * 1.5
+        # if np.isnan(max_logp): max_logp = 1
+        # logp.loc[Iloc] = max_logp
         # Create new adjmat based on indepdence test
         adjmat = vec2adjmat(source, target, weights=logp, symmetric=True, aggfunc='sum', verbose=verbose)
     else:
         adjmat = model['adjmat'].copy()
 
+    if edge_labels is not None:
+        edge_labels = list(adjmat.T.columns)
+
     # Make the dot and output Directed graph source code in the DOT language.
-    dot_graph = make_dot(adjmat.T.values.astype(float), labels=list(adjmat.T.columns), lower_limit=0, **params)
+    dot_graph = make_dot(adjmat.T.values.astype(float), labels=edge_labels, lower_limit=0, **params)
 
     # Return
     return dot_graph