Changed device of torch dataset in predict method to avoid potential …

…error
Project-Resilience · Jul 31, 2024 · c0e5300 · c0e5300
1 parent bef6351
commit c0e5300
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/prsdk/predictors/neural_network/neural_net_predictor.py b/prsdk/predictors/neural_network/neural_net_predictor.py
@@ -189,7 +189,7 @@ def predict(self, context_actions_df: pd.DataFrame) -> pd.DataFrame:
         :return: DataFrame of predictions properly labeled and indexed.
         """
         X_test_scaled = self.scaler.transform(context_actions_df[self.features])
-        test_ds = TorchDataset(X_test_scaled, np.zeros(len(X_test_scaled)))
+        test_ds = TorchDataset(X_test_scaled, np.zeros(len(X_test_scaled)), device=self.device)
         test_dl = DataLoader(test_ds, self.batch_size, shuffle=False)
         pred_list = []
         with torch.no_grad():