diff --git a/visualtorch/graph.py b/visualtorch/graph.py
index 181cbfd..f4513eb 100644
--- a/visualtorch/graph.py
+++ b/visualtorch/graph.py
@@ -168,7 +168,6 @@ def _draw_connector(
 
 def _retrieve_isbox_units(layer: torch.autograd.Function, show_neurons: bool) -> tuple[bool, int]:
     """Return the number of units and the flag whether to visualize using a box or not."""
-    print("test: ", type(layer))
     is_box = True
     units = 1
     if show_neurons: