Passing additional arguments in trainer.test()

I am trying to save the prediction results of top 3 checkpoints I have saved. I was wondering if there is a way to pass an argument in trainer.test(ckpt_path=top_{idx}.ckpt,idx) so that I could reference the number of K within the test_epoch_end().

here is my main function

topk_checkpoint_paths = os.listdir(checkpoint_path)
print(os.path.join(checkpoint_path, topk_checkpoint_paths[0]))
for i in range(train_params["save_top_k"]):
    trainer.test(ckpt_path=os.path.join(checkpoint_path, topk_checkpoint_paths[i]))

and what I want to do is to save a csv file in test_epoch_end with a reference to idx of top_k

        test_label_path = f"predictions/{self.dataset_params['experiment_name']}_{self.dataset_params['inference_save']}_with_label.csv"
        results_path = f"predictions/{self.dataset_params['experiment_name']}_{self.dataset_params['inference_save']}_results.csv"

I can solve this issue with making a class variable such as self.GLOBAL_ID but I want to know if there is a more “lightning” way of doing it.

I would appreciate any kind of help! thanks

hello :slight_smile:
Excise me but … I am not sure if you can save a model with the test mode :no_mouth:
Please, correct me of I am wrong!
saving a model need to be with checkpoint like this :

checkpoint_callback = ModelCheckpoint(
  save_top_k=1, # here you can save top 3 like you want
  monitor="Validation loss",