diff --git a/competition/random_forest.py b/competition/random_forest.py
index 0540f0a..7e6e5ed 100644
--- a/competition/random_forest.py
+++ b/competition/random_forest.py
@@ -20,7 +20,7 @@ df.drop(['flowStartMilliseconds'], 1, inplace=True)
 X = np.array(df.drop(columns=['sublabel']))
 y = np.array(df['sublabel'])
 
-X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
+X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y)
 
 clf = RandomForestClassifier(n_estimators=50, n_jobs=-1, criterion='gini', random_state=0)
 clf.fit(X_train, y_train)