Add emb_mode required argument into load_model

#302
Files changed (1) hide show
  1. geneformer/emb_extractor.py +1 -1
geneformer/emb_extractor.py CHANGED
@@ -565,7 +565,7 @@ class EmbExtractor:
565
  filtered_input_data, cell_state, self.nproc
566
  )
567
  downsampled_data = pu.downsample_and_sort(filtered_input_data, self.max_ncells)
568
- model = pu.load_model(self.model_type, self.num_classes, model_directory)
569
  layer_to_quant = pu.quant_layers(model) + self.emb_layer
570
  embs = get_embs(
571
  model,
 
565
  filtered_input_data, cell_state, self.nproc
566
  )
567
  downsampled_data = pu.downsample_and_sort(filtered_input_data, self.max_ncells)
568
+ model = pu.load_model(self.model_type, self.num_classes, model_directory, self.emb_mode)
569
  layer_to_quant = pu.quant_layers(model) + self.emb_layer
570
  embs = get_embs(
571
  model,