Thanks, this could work if I knew the indices of the chosen rows that are passed to fit!
The smaller matrix is built from choosing a subset of rows and columns and sampling their intersection. That is, if K is 10x10 and train_idx=[1,4,6], then the smaller matrix is the 3x3 K[train_idx,train_idx].
Hence, since I need the indices I wonder whether I should define a new sampler, or whatever function passes the data to train!, to obtain the smaller square matrix to be used in training