diff --git a/qlib/contrib/model/pytorch_alstm_ts.py b/qlib/contrib/model/pytorch_alstm_ts.py index 4cbd8a49e..41be0824d 100644 --- a/qlib/contrib/model/pytorch_alstm_ts.py +++ b/qlib/contrib/model/pytorch_alstm_ts.py @@ -262,7 +262,7 @@ class ALSTM(Model): dl_test = dataset.prepare("test", data_key=DataHandlerLP.DK_I) dl_test.config(fillna_type="ffill+bfill") - test_loader = DataLoader(dl_test, batch_size=self.batch_size, num_workers=20) + test_loader = DataLoader(dl_test, batch_size=self.batch_size, num_workers=self.n_jobs) self.ALSTM_model.eval() preds = [] diff --git a/qlib/contrib/model/pytorch_gats_ts.py b/qlib/contrib/model/pytorch_gats_ts.py index 534e5e27a..633da075a 100644 --- a/qlib/contrib/model/pytorch_gats_ts.py +++ b/qlib/contrib/model/pytorch_gats_ts.py @@ -307,7 +307,7 @@ class GATs(Model): dl_test = dataset.prepare("test", data_key=DataHandlerLP.DK_I) dl_test.config(fillna_type="ffill+bfill") - test_loader = DataLoader(dl_test, batch_size=self.batch_size, num_workers=20) + test_loader = DataLoader(dl_test, batch_size=self.batch_size, num_workers=self.n_jobs) self.ALSTM_model.eval() preds = [] diff --git a/qlib/contrib/model/pytorch_gru_ts.py b/qlib/contrib/model/pytorch_gru_ts.py index b1c943a89..149c9f8d0 100755 --- a/qlib/contrib/model/pytorch_gru_ts.py +++ b/qlib/contrib/model/pytorch_gru_ts.py @@ -262,7 +262,7 @@ class GRU(Model): dl_test = dataset.prepare("test", data_key=DataHandlerLP.DK_I) dl_test.config(fillna_type="ffill+bfill") - test_loader = DataLoader(dl_test, batch_size=self.batch_size, num_workers=20) + test_loader = DataLoader(dl_test, batch_size=self.batch_size, num_workers=self.n_jobs) self.GRU_model.eval() preds = [] diff --git a/qlib/contrib/model/pytorch_lstm_ts.py b/qlib/contrib/model/pytorch_lstm_ts.py index c717d3955..a7f8a2444 100755 --- a/qlib/contrib/model/pytorch_lstm_ts.py +++ b/qlib/contrib/model/pytorch_lstm_ts.py @@ -262,7 +262,7 @@ class LSTM(Model): dl_test = dataset.prepare("test", data_key=DataHandlerLP.DK_I) dl_test.config(fillna_type="ffill+bfill") - test_loader = DataLoader(dl_test, batch_size=self.batch_size, num_workers=20) + test_loader = DataLoader(dl_test, batch_size=self.batch_size, num_workers=self.n_jobs) self.LSTM_model.eval() preds = [] diff --git a/qlib/contrib/model/pytorch_sfm_ts.py b/qlib/contrib/model/pytorch_sfm_ts.py index 617c19b52..d4c149987 100644 --- a/qlib/contrib/model/pytorch_sfm_ts.py +++ b/qlib/contrib/model/pytorch_sfm_ts.py @@ -433,7 +433,7 @@ class SFM(Model): dl_test = dataset.prepare("test", data_key=DataHandlerLP.DK_I) dl_test.config(fillna_type="ffill+bfill") - test_loader = DataLoader(dl_test, batch_size=self.batch_size, num_workers=20) + test_loader = DataLoader(dl_test, batch_size=self.batch_size, num_workers=self.n_jobs) self.sfm_model.eval() preds = []