def _wrap_and_compile_model_for_evaluation()

in optimum/graphcore/trainer_seq2seq.py [0:0]

• 7 lines of code
• 7 McCabe index (conditional complexity)

    def _wrap_and_compile_model_for_evaluation(self, dataloader, prediction_loss_only):
        if prediction_loss_only:
            return super()._wrap_and_compile_model_for_evaluation(dataloader, prediction_loss_only)

        # Unwrap the model, including parameter and buffer annotations and the
        # model as a whole. This is needed to avoid issues with persistent buffers
        # when compiling an inference model for generation
        unwrapModelIfNecessary(self.model)

        # reparallelize for generation
        self.model.deparallelize().ipu_config.eval()
        self.model.parallelize(for_generation=True, **self.model.ipu_config.inference_parallelize_kwargs)

        # let IPUGenerationMixin::_call_generate handle compilation of the model
        # note though that self.model.poptorch_decoder and self.model.poptorch_encoder
        # (attribute added by IPUGenerationMixin::_call_generate)
        # are the actual models attached to the device
        return self.model