def _update_model_with_head_info()

in summarize_from_feedback/query_response_model.py [0:0]

• 12 lines of code
• 2 McCabe index (conditional complexity)

    def _update_model_with_head_info(self, model):
        if not self.logit_head:
            model.include_output_unembeddings = False
            model.unembedding_weights = None
        model = ModelWithHeads(
            model,
            scalar_heads=list(self.heads),
            d_model=model.d_model,
            init_scales=self.init_scales,
        )
        model = model.to(self.device)
        return model