Add handling empty scaler

Piotr45 · Jan 18, 2023 · 375316d · 375316d
1 parent b6a429d
commit 375316d
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/neural-network/scripts/prepare_dataset_pipeline.py b/neural-network/scripts/prepare_dataset_pipeline.py
@@ -53,13 +53,14 @@ def get_datasets(
     # all features are actually the same feature from different
     # timestamps
 
-    X_train_flattened = X_train.values.reshape(-1, 1)
-    scaler.fit(X_train_flattened)
-
-    for col in X_train.columns:
-        X_train[[col]] = scaler.transform(X_train[[col]])
-        X_val[[col]] = scaler.transform(X_val[[col]])
-        X_test[[col]] = scaler.transform(X_test[[col]])
+    if scaler is not None:
+        X_train_flattened = X_train.values.reshape(-1, 1)
+        scaler.fit(X_train_flattened)
+
+        for col in X_train.columns:
+            X_train[[col]] = scaler.transform(X_train[[col]])
+            X_val[[col]] = scaler.transform(X_val[[col]])
+            X_test[[col]] = scaler.transform(X_test[[col]])
 
     train_set = (
         tf.data.Dataset.from_tensor_slices((X_train, y_train))