Fix examples

polyaxon · Mar 19, 2022 · 15f671e · 15f671e
1 parent 5b08c38
commit 15f671e
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 6 deletions.
diff --git a/in_cluster/lightgbm/wine/model.py b/in_cluster/lightgbm/wine/model.py
@@ -66,7 +66,7 @@
     tracking.log_data_ref(content=X_train, name='x_train')
     tracking.log_data_ref(content=y_train, name='y_train')
     tracking.log_data_ref(content=X_test, name='X_test')
-    tracking.log_data_ref(content=y_test, name='y_train')
+    tracking.log_data_ref(content=y_test, name='y_test')
 
     lgb_train = lgb.Dataset(X_train, y_train)
     lgb_eval = lgb.Dataset(X_test, y_test, reference=lgb_train)

diff --git a/in_cluster/xgboost/boston/model.py b/in_cluster/xgboost/boston/model.py
@@ -71,7 +71,7 @@
     }
 
     # Polyaxon
-    tracking.init()
+    tracking.init(is_offline=True)
 
     boston = load_boston()
     data = pd.DataFrame(boston.data)
@@ -84,19 +84,20 @@
     tracking.log_data_ref(content=X_train, name='x_train')
     tracking.log_data_ref(content=y_train, name='y_train')
     tracking.log_data_ref(content=X_test, name='X_test')
-    tracking.log_data_ref(content=y_test, name='y_train')
+    tracking.log_data_ref(content=y_test, name='y_test')
+    callback = polyaxon_callback()
 
     dtrain = xgb.DMatrix(X_train, label=y_train)
     dtest = xgb.DMatrix(X_test, label=y_test)
 
     if args.cross_validate:
-        xgb.cv(params, dtrain, num_boost_round=20, nfold=7, callbacks=[polyaxon_callback()])
+        xgb.cv(params, dtrain, num_boost_round=20, nfold=7, callbacks=[callback])
     else:
         xgb.train(
             params,
             dtrain,
             20,
             [(dtest, 'eval'), (dtrain, 'train')],
-            callbacks=[polyaxon_callback()]  # Polyaxon
+            callbacks=[callback]
         )
 
diff --git a/tracking/xgboost/iris/model.py b/tracking/xgboost/iris/model.py
@@ -66,7 +66,7 @@ def model(log_learning_rate, max_depth=3, num_rounds=10, min_child_weight=5):
     experiment.log_data_ref(content=X_train, name='x_train')
     experiment.log_data_ref(content=y_train, name='y_train')
     experiment.log_data_ref(content=X_test, name='X_test')
-    experiment.log_data_ref(content=y_test, name='y_train')
+    experiment.log_data_ref(content=y_test, name='y_test')
 
     logger.info('Train model...')
     accuracy = model(log_learning_rate=args.log_learning_rate,