diff --git a/src/finn/transformation/code_gen_transformation.py b/src/finn/transformation/code_gen_transformation.py
index c015ca4ca6291e148410d9ca544ca904646e4ee2..6b0ab56ee826b6e2b296a7520ba725317c99c3ca 100644
--- a/src/finn/transformation/code_gen_transformation.py
+++ b/src/finn/transformation/code_gen_transformation.py
@@ -13,8 +13,7 @@ def code_gen_transformation(node, context, model):
         # lookup op_type in registry of CustomOps
         inst = registry.custom_op[op_type](node)
 
-        # get the path of the code generation directory if already set
-        # check instance and check node attributes for value
+        # get the path of the code generation directory
         code_gen_dir = inst.code_gen_dir
         code_gen_dir = code_gen_dir.s.decode("UTF-8")
 
@@ -50,6 +49,7 @@ def code_gen_transformation(node, context, model):
                             model.set_attribute(node, "code_gen_dir", tmp_dir)
                     else:
                         raise Exception("Code was not generated!")
+                # else: attribute is correctly set
             else:
                 inst.code_gen_dir = tmp_dir
                 model.set_attribute(node, "code_gen_dir", tmp_dir)
diff --git a/tests/transformation/test_code_gen_trafo.py b/tests/transformation/test_code_gen_trafo.py
index 4f050fc1f7738364fde8391fff5e0c8925835895..01768a36c8322b39ad1cde286fde008ab2b7c6d9 100644
--- a/tests/transformation/test_code_gen_trafo.py
+++ b/tests/transformation/test_code_gen_trafo.py
@@ -1,3 +1,5 @@
+import os
+
 import numpy as np
 from onnx import TensorProto, helper
 
@@ -59,3 +61,18 @@ def test_code_gen_trafo():
     model.set_initializer("thresh", T)
 
     model = model.transform(CodeGen())
+    for node in model.graph.node:
+        code_gen_attribute = util.get_by_name(node.attribute, "code_gen_dir")
+        tmp_dir = code_gen_attribute.s.decode("UTF-8")
+        assert os.path.isdir(
+            tmp_dir
+        ), """Code generation directory of node with
+            op type {} does not exist!""".format(
+            node.op_type
+        )
+        assert (
+            len(os.listdir(tmp_dir)) != 0
+        ), """Code generation directory of node with
+            op type {} is empty!""".format(
+            node.op_type
+        )