Mirrors/the-algorithm-ml
2
0
Fork 0
mirror of https://github.com/twitter/the-algorithm-ml.git synced 2024-12-22 22:31:48 +01:00
Code Issues Packages Projects Releases Wiki Activity

Early exit

Browse Source 
used the any function to short-circuit the loop as soon as an empty tensor is found.
This commit is contained in:
Darshan P 2023-04-03 16:00:13 +00:00
parent c4c5072402
commit 1c0d0499ff
1 changed files with 10 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
core/metric_mixin.py
View File

@ -34,17 +34,17 @@ import torchmetrics
class MetricMixin:
@abstractmethod
def transform(self, outputs: Dict[str, torch.Tensor]) -> Dict:
...
@abstractmethod
def transform(self, outputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
...
def update(self, outputs: Dict[str, torch.Tensor]):
results = self.transform(outputs)
# Do not try to update if any tensor is empty as a result of stratification.
if any((torch.is_tensor(value) and not value.nelement()) for value in results.values()):
return
super().update(**results)
def update(self, outputs: Dict[str, torch.Tensor]):
results = self.transform(outputs)
# Do not try to update if any tensor is empty as a result of stratification.
for value in results.values():
if torch.is_tensor(value) and not value.nelement():
return
super().update(**results)
class TaskMixin: