diff options
| author | Jordan Gong <jordan.gong@protonmail.com> | 2021-01-09 20:53:29 +0800 |
|---|---|---|
| committer | Jordan Gong <jordan.gong@protonmail.com> | 2021-01-09 20:53:29 +0800 |
| commit | 6f278a962d70e90ac530f5723e198c7c356e8297 (patch) | |
| tree | eef4663d708a6b48e6a2a07a40ebb0a5fb933fc5 /utils | |
| parent | dd12098603ac415904b9a8d512889deb995a8391 (diff) | |
Fix NaN when separate sum is zero
Diffstat (limited to 'utils')
| -rw-r--r-- | utils/triplet_loss.py | 5 |
1 files changed, 3 insertions, 2 deletions
diff --git a/utils/triplet_loss.py b/utils/triplet_loss.py index 1d63a0e..1899cc9 100644 --- a/utils/triplet_loss.py +++ b/utils/triplet_loss.py @@ -30,8 +30,9 @@ class BatchAllTripletLoss(nn.Module): all_loss = F.relu(self.margin + positive_negative_dist).view(p, -1) # Non-zero parted mean - parted_loss_mean = all_loss.sum(1) / (all_loss != 0).sum(1) - parted_loss_mean[parted_loss_mean == float('Inf')] = 0 + non_zero_counts = (all_loss != 0).sum(1) + parted_loss_mean = all_loss.sum(1) / non_zero_counts + parted_loss_mean[non_zero_counts == 0] = 0 loss = parted_loss_mean.sum() return loss |