fix: regression gpu integration an typos

dreamquark-ai · Dec 4, 2019 · 269b4c5 · 269b4c5
1 parent 3626d2f
commit 269b4c5
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/pytorch_tabnet/tab_model.py b/pytorch_tabnet/tab_model.py
@@ -502,7 +502,7 @@ def train_batch(self, data, targets):
  batch_outs = {'loss': loss_value,
  'y_preds': output,
  'y': targets,
- 'batch_importance': M_explain.sum(dim=0).detach().numpy()}
+ 'batch_importance': M_explain.sum(dim=0).cpu().detach().numpy()}
  return batch_outs
 
  def predict_epoch(self, loader):
@@ -774,7 +774,7 @@ def train_batch(self, data, targets):
  batch_outs = {'loss': loss_value,
  'y_preds': output,
  'y': targets,
- 'batch_importance': M_explain.sum(dim=0).detach().numpy()}
+ 'batch_importance': M_explain.sum(dim=0).cpu().detach().numpy()}
  return batch_outs
 
  def predict_epoch(self, loader):

diff --git a/regression_example.ipynb b/regression_example.ipynb
@@ -279,7 +279,7 @@
  " n_estimators=1000,\n",
  " verbosity=0,\n",
  " silent=None,\n",
- " objective='binary:logistic',\n",
+ " objective='reg:linear',\n",
  " booster='gbtree',\n",
  " n_jobs=-1,\n",
  " nthread=None,\n",
@@ -310,11 +310,11 @@
  "outputs": [],
  "source": [
  "preds = np.array(clf_xgb.predict(X_valid))\n",
- "valid_auc = mean_squared_error(y_pre=preds, y_true=y_valid)\n",
+ "valid_auc = mean_squared_error(y_pred=preds, y_true=y_valid)\n",
  "print(valid_auc)\n",
  "\n",
  "preds = np.array(clf_xgb.predict(X_test))\n",
- "test_auc = mean_squared_error(y_pre=preds, y_true=y_test)\n",
+ "test_auc = mean_squared_error(y_pred=preds, y_true=y_test)\n",
  "print(test_auc)"
  ]
  }