summaryrefslogtreecommitdiff
path: root/models/rgb_part_net.py
diff options
context:
space:
mode:
authorJordan Gong <jordan.gong@protonmail.com>2021-01-06 21:26:56 +0800
committerJordan Gong <jordan.gong@protonmail.com>2021-01-06 21:26:56 +0800
commit4befe59046fb3adf8ef8eb589999a74cf7136ff6 (patch)
tree10a4b47c6dbbd834138beb8bb57c1e19d7cac068 /models/rgb_part_net.py
parentd473b9d1d79bf185b1811ce403f82fdd68fb366c (diff)
Add TensorBoard support
Diffstat (limited to 'models/rgb_part_net.py')
-rw-r--r--models/rgb_part_net.py2
1 files changed, 1 insertions, 1 deletions
diff --git a/models/rgb_part_net.py b/models/rgb_part_net.py
index a58be39..73d5952 100644
--- a/models/rgb_part_net.py
+++ b/models/rgb_part_net.py
@@ -78,7 +78,7 @@ class RGBPartNet(nn.Module):
batch_all_triplet_loss = self.ba_triplet_loss(x, y)
losses = (*losses, batch_all_triplet_loss)
loss = torch.sum(torch.stack(losses))
- return loss, (loss.item() for loss in losses)
+ return loss, [loss.item() for loss in losses]
else:
return x