make MLP compatible with the new pyhealth 2.0 loaders (#543)

Co-authored-by: John Wu <[email protected]>

Author: jhnwu3

pyhealth/models/embedding.py

@@ -4,14 +4,14 @@
 import torch.nn as nn
 
 from ..datasets import SampleDataset
-from ..processors import SequenceProcessor, TimeseriesProcessor
+from ..processors import SequenceProcessor, TimeseriesProcessor, TensorProcessor
 from .base_model import BaseModel
 
 
 class EmbeddingModel(BaseModel):
     """
     EmbeddingModel is responsible for creating embedding layers for different types of input data.
-    
+
     Attributes:
         dataset (SampleDataset): The dataset containing input processors.
         embedding_layers (nn.ModuleDict): A dictionary of embedding layers for each input field.
@@ -33,13 +33,27 @@ def __init__(self, dataset: SampleDataset, embedding_dim: int = 128):
                 self.embedding_layers[field_name] = nn.Embedding(
                     num_embeddings=vocab_size,
                     embedding_dim=embedding_dim,
-                    padding_idx=0
+                    padding_idx=0,
                 )
             elif isinstance(processor, TimeseriesProcessor):
                 self.embedding_layers[field_name] = nn.Linear(
-                    in_features=processor.size,
-                    out_features=embedding_dim
+                    in_features=processor.size, out_features=embedding_dim
                 )
+            elif isinstance(processor, TensorProcessor):
+                # For tensor processor, we need to determine the input size
+                # from the first sample in the dataset
+                sample_tensor = None
+                for sample in dataset.samples:
+                    if field_name in sample:
+                        sample_tensor = processor.process(sample[field_name])
+                        break
+                if sample_tensor is not None:
+                    input_size = (
+                        sample_tensor.shape[-1] if sample_tensor.dim() > 0 else 1
+                    )
+                    self.embedding_layers[field_name] = nn.Linear(
+                        in_features=input_size, out_features=embedding_dim
+                    )
 
     def forward(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
         """