def METHOD_NAME(self, fl_ctx: FLContext):
    (self.train_images, self.train_labels), (
        self.test_images,
        self.test_labels,
    ) = tf.keras.datasets.mnist.load_data()
    self.train_images, self.test_images = (
        self.train_images / 255.0,
        self.test_images / 255.0,
    )
    # simulate separate datasets for each client by dividing MNIST dataset in half
    client_name = fl_ctx.get_identity_name()
    if client_name == "site-1":
        self.train_images = self.train_images[: len(self.train_images) // 2]
        self.train_labels = self.train_labels[: len(self.train_labels) // 2]
        self.test_images = self.test_images[: len(self.test_images) // 2]
        self.test_labels = self.test_labels[: len(self.test_labels) // 2]
    elif client_name == "site-2":
        self.train_images = self.train_images[len(self.train_images) // 2 :]
        self.train_labels = self.train_labels[len(self.train_labels) // 2 :]
        self.test_images = self.test_images[len(self.test_images) // 2 :]
        self.test_labels = self.test_labels[len(self.test_labels) // 2 :]
    model = Net()
    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    model.compile(optimizer="adam", loss=loss_fn, metrics=["accuracy"])
    _ = model(tf.keras.Input(shape=(28, 28)))
    self.var_list = [model.get_layer(index=index).name for index in range(len(model.get_weights()))]
    self.model = model