diff --git a/pointnet2/model.py b/pointnet2/model.py
index a67e007..a63dba3 100644
--- a/pointnet2/model.py
+++ b/pointnet2/model.py
@@ -134,7 +134,7 @@ def __init__(
             raise ValueError(f"Received unknown feature norm `{feature_norm}`!")
 
         self.segmentation_layers = [
-            FeaturePropagationModule(mlp_units=units, feature_norm="batch")
+            FeaturePropagationModule(mlp_units=units, feature_norm=feature_norm)
             for units in fp_units
         ]