make resnet and mlp responsible to ensure embeddings are correct dims

OHDSI · Nov 25, 2024 · 62bd2ba · 62bd2ba
1 parent ed51ca5
commit 62bd2ba
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/inst/python/MultiLayerPerceptron.py b/inst/python/MultiLayerPerceptron.py
@@ -53,6 +53,8 @@ def __init__(
     def forward(self, input):
         x_cat = input["cat"]
         x_cat = self.embedding(x_cat)
+        if x_cat.dim() == 3:
+            x_cat = x_cat.mean(dim=1)
         if "num" in input.keys() and self.num_embedding is not None:
             x_num = input["num"]
             x = (x_cat + self.num_embedding(x_num).mean(dim=1)) / 2

diff --git a/inst/python/ResNet.py b/inst/python/ResNet.py
@@ -66,6 +66,8 @@ def __init__(
     def forward(self, x):
         x_cat = x["cat"]
         x_cat = self.embedding(x_cat)
+        if x_cat.dim() == 3:
+            x_cat = x_cat.mean(dim=1)
         if (
             "num" in x.keys()
             and x["num"] is not None