(self, model, device)
| 304 | |
| 305 | class Shared_grad_buffers(): |
| 306 | def __init__(self, model, device): |
| 307 | self.grads = {} |
| 308 | self.counter = Counter() |
| 309 | |
| 310 | for name, p in model.named_parameters(): |
| 311 | if p.requires_grad: |
| 312 | self.grads[name + '_grad'] = torch.zeros(p.size()).share_memory_().to(device) |
| 313 | |
| 314 | def add_gradient(self, model): |
| 315 | self.counter.increment() |