Args: all_grads: K x N, K lists of gradients to be packed
(self, all_grads, devices)
| 444 | return grads |
| 445 | |
| 446 | def pack_all(self, all_grads, devices): |
| 447 | """ |
| 448 | Args: |
| 449 | all_grads: K x N, K lists of gradients to be packed |
| 450 | """ |
| 451 | ret = [] # #GPU x #split |
| 452 | for dev, grads in zip(devices, all_grads): |
| 453 | with tf.device(dev): |
| 454 | ret.append(self.pack(grads)) |
| 455 | return ret |
| 456 | |
| 457 | def unpack_all(self, all_packed, devices): |
| 458 | """ |