fix: functional balanced version

dreamquark-ai · Nov 20, 2019 · fab7f16 · fab7f16
1 parent 0bf45d2
commit fab7f16
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/pytorch_tabnet/tab_model.py b/pytorch_tabnet/tab_model.py
@@ -167,11 +167,9 @@ def fit(self, X_train, y_train, X_valid=None, y_valid=None,
  sampler = WeightedRandomSampler(samples_weight, len(samples_weight))
  train_dataloader = DataLoader(TorchDataset(X_train, y_train),
  batch_size=self.batch_size, sampler=sampler)
- valid_dataloader = DataLoader(TorchDataset(X_valid, y_valid),
- batch_size=self.batch_size, shuffle=False)
-
- train_dataloader = DataLoader(TorchDataset(X_train, y_train),
- batch_size=self.batch_size, shuffle=True)
+ else:
+ train_dataloader = DataLoader(TorchDataset(X_train, y_train),
+ batch_size=self.batch_size, shuffle=True)
  valid_dataloader = DataLoader(TorchDataset(X_valid, y_valid),
  batch_size=self.batch_size, shuffle=False)