Possible Bug in Remove Single Ts TimeRankFromBatch

RLlib
1

I ran into an issue and I think that it is trying to squeeze when there is no dimension in the data.

in RemoveSingleTsTimeRankFromBatch, it might have to be changed to,

Only squeeze if the item has more than one dimension and the first dimension is 1

if mid is None or rl_module[mid].is_stateful():
    if isinstance(item, np.ndarray) and len(item.shape) > 1 and item.shape[0] == 1:
        return np.squeeze(item, axis=0)
return item