def METHOD_NAME(self):
    name = rpc.get_worker_info().name
    m = self.ps_rref.rpc_sync().get_model()
    for inputs, labels in self.get_next_batch():
        timed_log(f"{name} processing one batch")
        self.loss_fn(m(inputs), labels).backward()
        timed_log(f"{name} reporting grads")
        m = rpc.rpc_sync(
            self.ps_rref.owner(),
            BatchUpdateParameterServer.update_and_fetch_model,
            args=(self.ps_rref, [p.grad for p in m.cpu().parameters()]),
        )
        timed_log(f"{name} got updated model")