Method
tock
(self, name="default", evict=True, print_ms=True)
Source from the content-addressed store, hash-verified
| 78 | ms = self.tock(name, evict=True, print_ms=False) |
| 79 | |
| 80 | def tock(self, name="default", evict=True, print_ms=True): |
| 81 | if name in self.ends: |
| 82 | self.ends[name].record() |
| 83 | torch.cuda.synchronize() |
| 84 | ms = self.starts[name].elapsed_time(self.ends[name]) |
| 85 | if name not in self.agg: |
| 86 | self.agg[name] = 0.0 |
| 87 | self.agg[name] += ms |
| 88 | if evict: |
| 89 | self.starts.pop(name) |
| 90 | self.ends.pop(name) |
| 91 | |
| 92 | if print_ms and name in self.agg: |
| 93 | print(f"{name} took: {self.agg[name] / 1000.0:.5f}s") |
| 94 | |
| 95 | return self.agg[name] |
| 96 | |
| 97 | def reset(self): |
| 98 | self.starts = {} |
Tested by
no test coverage detected