Further update:
If I set train_idx = 100 so that the shape of training and testing data is equal, then the error does not occur.
Does this mean I am misinterpreting how I should be using set_data, or could this be a bug where the size is just being passed on from the training data set?