Skip to content

Commit e62684e

Browse files
fix bug on TRA dataset (#1135)
* fix bug on TRA dataset solve issue "qrun TRA model error (#1062)" * apply black pylint
1 parent 8a5efda commit e62684e

File tree

1 file changed

+8
-2
lines changed

1 file changed

+8
-2
lines changed

qlib/contrib/data/dataset.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -203,8 +203,14 @@ def setup_data(self, handler_kwargs: dict = None, **kwargs):
203203

204204
def _prepare_seg(self, slc, **kwargs):
205205
fn = _get_date_parse_fn(self._index[0][1])
206-
start_date = fn(slc.start)
207-
end_date = fn(slc.stop)
206+
if isinstance(slc, slice):
207+
start, stop = slc.start, slc.stop
208+
elif isinstance(slc, (list, tuple)):
209+
start, stop = slc
210+
else:
211+
raise NotImplementedError(f"This type of input is not supported")
212+
start_date = pd.Timestamp(fn(start))
213+
end_date = pd.Timestamp(fn(stop))
208214
obj = copy.copy(self) # shallow copy
209215
# NOTE: Seriable will disable copy `self._data` so we manually assign them here
210216
obj._data = self._data # reference (no copy)

0 commit comments

Comments
 (0)