Updates usage of get_specs API to accept transformer_config instead o…

…f transformer_config.num_moe_experts NVIDIA/NeMo#9035
NVIDIA · Sep 16, 2024 · 629c862 · 629c862
1 parent 745d148
commit 629c862
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo_aligner/models/nlp/gpt/megatron_gpt_reward_model.py b/nemo_aligner/models/nlp/gpt/megatron_gpt_reward_model.py
@@ -92,7 +92,7 @@ def model_provider_func(self, pre_process, post_process):
 
         model = GPTRewardModel(
             config=self.transformer_config,
-            transformer_layer_spec=get_specs(self.spec_name, self.transformer_config.num_moe_experts),
+            transformer_layer_spec=get_specs(self.spec_name, self.transformer_config),
             vocab_size=self.cfg.get("override_vocab_size", self.padded_vocab_size),
             max_sequence_length=self.cfg.get("encoder_seq_length", 512),
             pre_process=pre_process,