r"""Mark send/recv need gradiants by key. Args: key: key to match send/recv op. is_grad: whether this op need grad.
(self, key, is_grad)
| 40 | return self.mm_server_port |
| 41 | |
| 42 | def set_is_grad(self, key, is_grad): |
| 43 | r"""Mark send/recv need gradiants by key. |
| 44 | |
| 45 | Args: |
| 46 | key: key to match send/recv op. |
| 47 | is_grad: whether this op need grad. |
| 48 | """ |
| 49 | with self.lock: |
| 50 | future = self.dict_is_grad[key] |
| 51 | future.set(is_grad) |
| 52 | return True |
| 53 | |
| 54 | def check_is_grad(self, key): |
| 55 | r"""Check whether send/recv need gradiants. |