Can MatchFinder/InferenceModel return distances for all classes?

[vemchance]

As always, thanks again for the library and apologies for asking another question!

I was wondering if the MatchFinder or InferenceModel can return all of the distances for all classes or if it is expected that it will only return distances of 'matches'. I am using a late fusion technique which takes the distance of a query image to all classes and uses dynamic weighting to produce a more accurate prediction. I train two seperate modalities, one which uses the text and one only the images, to produce two different models, and use the distance of both to dynamically weight each modality for a query image.

Using class centers in this discussion, the dynamic weighting works extremely well and I can use the InferenceModel to return all the neighbours (e.g., a distance for the query image to all classes in the dataset). However, the accuracy of the individual modalities is different from in training (it is using the same dataset), though I expect this behaviour since I did not use the same class centre method in training. This code below returns distances for all classes:

class_centers = []
for L in torch.unique(labels, sorted=True):
    mask = labels == L
    i, l = np.where(mask.cpu())
    class_embeddings = embeddings[i, ...]
    center = torch.mean(class_embeddings, dim=0)
    class_centers.append(center)

class_centers = torch.stack(class_centers, dim=0)
knn = FaissKNN(reset_before=False, reset_after=False)
knn.train(class_centers)
inference_model = InferenceModel(model, knn_func=knn, data_device=device)

dataloader = torch.utils.data.DataLoader(val_dataset, batch_size = 1)

for data, labels in dataloader:
    data, labels = move_data_to_device(data, labels, device=device)
    distances, indices = inference_model.get_nearest_neighbors(data, k=classes)

MatchFinder will provide me with the same accuracy as in training. However, it will only return a certain number of classes and their distances, rather than the distances between the query and all classes, even if I set the number of neighbours to the total number of classes, modify the MatchFinder threshold, or remove the MatchFinder entirely. I'm not sure if this is expected behaviour but wanted to check or perhaps I'm misunderstanding the InferenceModel?

For example:

knn = FaissKNN(reset_before=False, reset_after=False)
match_finder = MatchFinder(distance=CosineSimilarity(), threshold=0.8) # modifying threshold doesn't return all classes
inference_model = InferenceModel(model, data_device=device, knn_func=knn, match_finder=match_finder) # removing matchfinder doesn't result all classes
inference_model.train_knn(train_dataset)

dataloader = torch.utils.data.DataLoader(val_dataset,  batch_size = 1) 

for data, labels in dataloader:
    data, labels = move_data_to_device(data, labels, device=device)

    # Get distances and indices for the nearest classes
    distances, indices = inference_model.get_nearest_neighbors(data, k=classes) # k is set to the total number of classes

The above code will return only a few distances (32 as opposed to 213, which is the total number of classes) even when adjusting or removing the MatchFinder. Usually, I get between 20 - 32 distances rather than 213.

I would like to use the InferenceModel as it is above instead of calculating the class centers, as I can compare my fusion method to my original results. It seems the InferenceModel method gives comparable results to training which I expected, but the fusion method requires all classes to have a distance calculated which I can't seem to retrieve even when increasing the number of neighbours.