def get_expected_model()

in modules/pipeline/custom_steps.py [0:0]

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

    def get_expected_model(self, model_name=None):
        """
        Build Sagemaker model representation of the expected trained model from
        the Training step. This can be passed to the ModelStep to save the
        trained model in Sagemaker.

        Args:
            model_name (str, optional): Specify a model name. If not provided,
            training job name will be used as the model name.
        Returns:
            sagemaker.model.Model: Sagemaker model representation of the
            expected trained model.
        """
        model = self.estimator.create_model()
        if model_name:
            model.name = model_name
        else:
            model.name = self.job_name
        model.model_data = self.output()["ModelArtifacts"]["S3ModelArtifacts"]
        return model