fix rixwew#21 add training to the F.dropout argument

wuyuup · Jul 27, 2020 · 8ac848d · 8ac848d
1 parent acc6997
commit 8ac848d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchfm/layer.py b/torchfm/layer.py
@@ -194,9 +194,9 @@ def forward(self, x):
         inner_product = p * q
         attn_scores = F.relu(self.attention(inner_product))
         attn_scores = F.softmax(self.projection(attn_scores), dim=1)
-        attn_scores = F.dropout(attn_scores, p=self.dropouts[0])
+        attn_scores = F.dropout(attn_scores, p=self.dropouts[0], training=self.training)
         attn_output = torch.sum(attn_scores * inner_product, dim=1)
-        attn_output = F.dropout(attn_output, p=self.dropouts[1])
+        attn_output = F.dropout(attn_output, p=self.dropouts[1], training=self.training)
         return self.fc(attn_output)