Skip to content

GitLab

"git@developer.sourcefind.cn:OpenDAS/deepspeed.git" did not exist on "7435b2f10af773b0204e77c3549b2b7df9a7a65b"
Unverified Commit 68c0cfbb authored by Peiqi Yin's avatar Peiqi Yin Committed by GitHub
Browse files

[Model] Fix diffpool loss (#3233)

parent 31772b14
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
examples/pytorch/diffpool/model/encoder.py
View file @ 68c0cfbb
......@@ -214,6 +214,8 @@ class DiffPool(nn.Module):
#softmax + CE
criterion = nn.CrossEntropyLoss()
loss = criterion(pred, label)
for key, value in self.first_diffpool_layer.loss_log.items():
loss += value
for diffpool_layer in self.diffpool_layers:
for key, value in diffpool_layer.loss_log.items():
loss += value
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment