| 807 | |
| 808 | |
| 809 | class _RemoteRecv(Function): |
| 810 | def __init__(self, op: RemoteRecv): |
| 811 | self.op = op |
| 812 | |
| 813 | def forward(self, dummy): |
| 814 | return apply(self.op, dummy)[0] |
| 815 | |
| 816 | def backward(self, grad): |
| 817 | get_client().bcast_val(grad is not None, self.op.key, 2) |
| 818 | if grad is not None: |
| 819 | remote_send(grad, self.op.rank_from) |
| 820 | |
| 821 | |
| 822 | def remote_send(inp: Tensor, dest_rank: int): |