How to use GPU for Keras learning in this example

I have the following code, for vertical federated learning where i use keras to create a simple 2 layer NN

batch_size = 32
learning_rate = 1e-3
epochs = 50

optimizer = keras.optimizers.Adam(learning_rate = learning_rate)
loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits = False)
train_acc_metric = tf.keras.metrics.SparseCategoricalAccuracy()

class Client():
    def __init__(self, train, test, labelled):
        self.__trainX = train.copy()
        self.__testX = test.copy()
        self.labelled = labelled
        
        if (labelled):
            self.__trainY = self.__trainX.pop('credit_risk')
            self.__testY = self.__testX.pop('credit_risk')
        
        normalizer = preprocessing.Normalization()
        normalizer.adapt(np.array(self.__trainX.loc[common_train_id]))
        
        self.model = keras.Sequential([
              normalizer,
              layers.Dense(128, activation='elu', kernel_regularizer=regularizers.l2(0.01)),
              layers.Dropout(0.5),
              layers.Dense(128, activation='elu', kernel_regularizer=regularizers.l2(0.01)),
              layers.Dropout(0.5),
              layers.Dense(2),
              layers.Softmax()])
        
    def next_batch(self, index):
        self.batchX = self.__trainX.loc[index]
        if not self.labelled:
            grads = []
            self.model_output = np.zeros((len(index), 2))
            for i in range(len(index)):
                with tf.GradientTape() as gt:
                    gt.watch(self.model.trainable_weights)
                    output_by_example = self.model(self.batchX.iloc[i:i+1], training = True)
                    output_for_grad = output_by_example[:,1]
                self.model_output[i] = output_by_example
                grads.append(gt.gradient(output_for_grad, self.model.trainable_weights))
            return grads
        else:
            self.batchY = self.__trainY.loc[index]
            with tf.GradientTape() as self.gt:
                self.gt.watch(self.model.trainable_weights)
                self.model_output = self.model(self.batchX, training = True)

    def cal_model(self):
        return self.model_output
  
    def predict(self, test_index):
        return self.model.predict(self.__testX.loc[test_index])# + 1e-8

    def test_answers(self, test_index):
        if self.labelled:
            return self.__testY.loc[test_index]
  
    def batch_answers(self):
        if self.labelled:
            return self.batchY

    def loss_and_update(self, a):
        if not self.labelled:
            raise AssertionError(This method can only be called by client 2)
        self.prob = (a + self.model_output)/2
        self.c = self.coefficient_and_update()/len(self.batchX)
        return self.prob, loss_fn(self.batchY, self.prob)
  
    def coefficient_and_update(self):
        if not self.labelled:
            raise AssertionError(This method can only be called by client 2)
        p = self.prob[:,1]
        c = (p-self.batchY)/((p)*(1-p))
        with self.gt:
            output = sum(c * self.model_output[:,1])/len(c)
        grads = self.gt.gradient(output, self.model.trainable_weights)
        optimizer.apply_gradients(zip(grads, self.model.trainable_weights))
        return c
  
    def update_with(self, grads):
        optimizer.apply_gradients(zip(grads, self.model.trainable_weights))

    def assemble_grad(self, partial_grads):
        if not self.labelled:
            raise AssertionError(This method can only be called by client 2)
    # to assemble the gradient for client 1
        for i in range(len(self.c)):
            partial_grads[i] = [x * self.c[i] for x in partial_grads[i]]
        return [sum(x) for x in zip(*partial_grads)]

client1 = Client(df_1_train, df_1_test, False)
client2 = Client(df_2_train, df_2_test, True)

epoch_loss = []
epoch_acc = []

for epoch in range(epochs):
    random.shuffle(common_train_id_list)
    train_id_batches = [common_train_id_list[i:i + batch_size] for i in range(0, len(common_train_id_list), batch_size)] 
    total_loss = 0.0
    # Iterate over the batches of the dataset.
    for step, batch_index in enumerate(train_id_batches):
    
        partial_grads = client1.next_batch(batch_index)
        client2.next_batch(batch_index)

        prob, loss_value = client2.loss_and_update(client1.cal_model())
        grad = client2.assemble_grad(partial_grads)
        client1.update_with(grad)

        total_loss = loss_value + total_loss
        train_acc_metric.update_state(client2.batch_answers(), prob)
  
    train_acc = train_acc_metric.result()
    train_acc_metric.reset_states()
    epoch_loss.append((total_loss)/(step + 1))
    epoch_acc.append(train_acc)

This is training on my CPU, and takes a long time. I want to train it on my GPU instead. How do i update the code to be able to run on gpu?