diff --git a/skais/ais/ais_trajectory.py b/skais/ais/ais_trajectory.py
index a20abf4c22039c35f570e22ac912c57dfa95f749..f4ec956b6930cca43121640fc4e2f93a698ce915 100644
--- a/skais/ais/ais_trajectory.py
+++ b/skais/ais/ais_trajectory.py
@@ -306,7 +306,6 @@ class AISTrajectory(AISPoints):
         else:
             bounds = []
             if type(features) is list:
-                nb_channels = len(features)
                 features_vectors = self.df[features].to_numpy()
                 for c in features_vectors.T:
                     bounds.append((0, max(c)))
@@ -319,6 +318,8 @@ class AISTrajectory(AISPoints):
                 features_vectors = self.df[features.keys()].to_numpy()
             else:
                 raise TypeError("Type not supported")
+
+            nb_channels = len(features_vectors.T)
             data = np.zeros((height, width, nb_channels), dtype=np.float)
 
 
diff --git a/skais/tests/ais/test_ais_trajectory.py b/skais/tests/ais/test_ais_trajectory.py
index d22e22f2c014b0f493ddbe05ab1270c8a41b6f0d..51d776ccfe534ef89e1784ad5936dd68b8695572 100644
--- a/skais/tests/ais/test_ais_trajectory.py
+++ b/skais/tests/ais/test_ais_trajectory.py
@@ -631,7 +631,7 @@ class TestAISTrajectoryImageGeneration(unittest.TestCase):
         )
 
         result = trajectory.generate_array_from_positions(height=9, width=18, link=True, bounding_box='fit',
-                                                          features={'sog': (0, 40), 'cog': (0, 40)}, node_size=2)
+                                                          features={'sog': (0, 40), 'cog': (0, 40)}, node_size=0)
         expected = np.array([[[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0],
                               [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0.5, 0.25]],
                              [[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [0, 0],