def test_data_splitter(batch_data, pred_len):
    """
    Split data [batch_size, total_len, 2] into datax and datay in train mode
    :param batch_data: data to be split
    :param pred_len: length of trajectories in final loss calculation
    :return: datax, datay
    """
    return batch_data[:, :-pred_len, :], batch_data[:, -pred_len:, :]