docs: def train_image(self):
          """Return the Docker image to use for training.
  
          The :meth:`~sagemaker.estimator.EstimatorBase.fit` method, which does the model training,
          calls this method to find the image to use for model training.
  
          Returns:
              str: The URI of the Docker image.
          """
          if self.image_name:
              return self.image_name
          else:
              return fw_utils.create_image_uri(self.sagemaker_session.boto_region_name,
                                               self._image_framework(),
                                               self.train_instance_type,
                                               self._image_version(),
                                               py_version=PYTHON_VERSION)