KeyError when running predict.py

[pailozka]

Hi,
it seems that parameter augment_genes does not propagate well. KeyError occurs when I'm running predict.py:

File "../shepherd/gene_prioritization_model.py", line 304, in inference
if self.hparams.hparams['augment_genes']:
KeyError: 'augment_genes'

[pailozka]

I had augment_genes set to False in hparams.py, so changing value extraction in gene_prioritization_model.py to

if self.hparams.get('augment_genes', False):

solves the problem.
The main problem stems from PyTorch Lightning functionality, which loads the model from the checkpoint and restores not only the weights but also the hyperparameters saved during training. This hparams dictionary for the checkpoint you provided may not contain augment_genes.
N.B. I'm running the following command:

python predict.py --run_type causal_gene_discovery --patient_data disease_simulated --edgelist KG_edgelist_mask.txt --node_map KG_node_map.txt --saved_node_embeddings_path checkpoints/pretrain.ckpt --best_ckpt checkpoints/causal_gene_discovery.ckpt