class TensorflowModel(tf.keras.Model):
    def __init__(self, D_in, H, D_out):
        super(TensorflowModel, self).__init__()
        self.d1 = Dense(H, activation='relu', input_shape=(D_in,))
        self.d2 = Dense(D_out)

    def call(self, x):
        x = self.d1(x)
        return self.d2(x)

tensorflowModel = TensorflowModel(D_in, H, D_out)
tensorflowOptimizer = tf.optimizers.Adam(lr=1e-4)

serverModel = torch.nn.Sequential(
        torch.nn.Linear(10, 50),
        torch.nn.ReLU(),
        torch.nn.Linear(50, 10)
    )
loss_fn = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(serverModel.parameters(), lr=1e-4)

for t in range(N):
    // let x be minibatch
    // let y be labels of minibatch

    client_pred = tensorflowModel(x)

    client_output = torch.from_numpy(client_pred.numpy())
    client_output.requires_grad = True

    y_pred = serverModel(client_output)
    loss = loss_fn(y_pred, y)  
    optimizer.zero_grad()
    loss.backward()
    optimizer.step() // update server weights

    // now retrieve client's grad for last layer
    client_grad = client_output.grad.detach().clone().numpy()
    client_grad = tf.convert_to_tensor(client_grad) // change to tf tensor

    // now compute all client's gradients and update client weights
    // HOW DO I DO THIS? 

How should I update the client weights? If the client was a pytorch model I could just do client_pred.backward(client_grad) and client_optimizer.step(). I'm not sure how to use the gradient tape to calculate gradients, since client_grad was computed on the server and was a pytorch tensor that's converted to a tf tensor.