Updating

app.py

@@ -14,8 +14,8 @@ def main(input1):
   train_set, test_set = train_test_split(housing, test_size=0.2, random_state=10)
 
   train_set_clean = train_set.dropna()
-  train_labels = train_set_clean["housing_median_age"].copy() # get labels for output label Y
-  train_features = train_set_clean.drop("housing_median_age", axis=1) # drop labels to get features X for training set
+  train_labels = train_set_clean[input1].copy() # get labels for output label Y
+  train_features = train_set_clean.drop(input1, axis=1) # drop labels to get features X for training set
 
   scaler = MinMaxScaler() ## define the transformer
   scaler.fit(train_features) ## call .fit() method to calculate the min and max value for each column in dataset