import tensorflow as tf

    def build_training_model(self, x, y):
        """
        This method will be called once by all data owners
        to create a local gradient computation on their machine.
        """
        _, _, grads = self._build_model(x, y)
        return grads

    def _build_validation_model(self, x, y):
        predictions, loss, _ = self._build_model(x, y)
        most_likely = tf.argmax(predictions, axis=1)
        return most_likely, loss

    def _build_data_pipeline(self):

        def normalize(image, label):
            image = tf.cast(image, tf.float32) / 255.0
            return image, label

        dataset = tf.data.TFRecordDataset(["./data/train.tfrecord"])
