How are you implementing the distance loss?
You could try to index the model’s output using the test class indices and compute your distance loss.
output = model(data)
output_test = output[:, test_indices]

However, I’m not sure, how your distance is calculated.
Also, I’m not familiar with your use case, but be careful using the testing data for training.