CHANGE: change calculation to .einsum(...)

This should be much more better to show the calculation logic inside the model.
p768lwy3 · Oct 29, 2019 · 0787dd5 · 0787dd5
1 parent 6be6057
commit 0787dd5
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/torecsys/layers/ctr/compose_excitation_network.py b/torecsys/layers/ctr/compose_excitation_network.py
@@ -89,6 +89,10 @@ def forward(self, field_emb_inputs: torch.Tensor) -> torch.Tensor:
         # inputs: field_emb_inputs, shape = (B, N * N, E)
         # output: outputs, shape = (B, N * N, E)
         attn_w = attn_w.unflatten("N", (("N", attn_w.size("N")), ("E", 1)))
-        outputs = field_emb_inputs * attn_w
+
+        # Multiply attentional weights on field embedding tensors
+        ## outputs = field_emb_inputs * attn_w
+        outputs = torch.einsum("ijk,ijh->ijk", [field_emb_inputs.rename(None), attn_w.rename(None)])
+        outputs.names = ("B", "N", "E")
 
         return outputs