I did some looking into this. Adding a lock on ParameterServer.forward has no effect.
    def forward(self, inp):
        # forward_lock defined globally
        with forward_lock:
            inp = inp.to(self.input_device)
            out = self.model(inp)
            # This output is forwarded over RPC, which as of 1.5.0 only accepts CPU tensors.
            # Tensors must be moved in and out of GPU memory due to this.
            out = out.to("cpu")
            return out