(self, device)
| 126 | |
| 127 | class SFTMetric: |
| 128 | def __init__(self, device): |
| 129 | self.n_step = 0 |
| 130 | self.right = torch.Tensor([0]).to(device=device) |
| 131 | self.total = torch.Tensor([0]).to(device=device) |
| 132 | self.total_loss = torch.Tensor([0]).to(device=device) |
| 133 | self.world_size = dist.get_world_size() |
| 134 | |
| 135 | def __call__(self, logits, labels, loss): |
| 136 | return self.update(logits, labels, loss) |