Repository: sunlabuiuc/PyHealth (language: python)

Issue:
Support new 6 EHR models
add 6 models (TCN, AdaCare, ConCare, GRASP, StageNet, Agent); update readme and docs


Discussion / hints:

Fix patch:
diff --git a/README.rst b/README.rst
index 1c4f42a9..a4de6e88 100644
--- a/README.rst
+++ b/README.rst
@@ -375,21 +375,27 @@ ISRUC                ``pyhealth.datasets.ISRUCDataset``       2016
 8. Machine/Deep Learning Models and Benchmarks :airplane:
 ----------------------------------------------------------
 
-==================================    ================  =================================  ======  ============================================================================================================================================================================
-Model Name                            Type              Module                             Year    Reference
-==================================    ================  =================================  ======  ============================================================================================================================================================================
-Multi-layer Perceptron                deep learning     ``pyhealth.models.MLP``            1986    `Backpropagation: theory, architectures, and applications <https://www.taylorfrancis.com/books/mono/10.4324/9780203763247/backpropagation-yves-chauvin-david-rumelhart>`_
-Convolutional Neural Network (CNN)    deep learning     ``pyhealth.models.CNN``            1989    `Handwritten Digit Recognition with a Back-Propagation Network <https://proceedings.neurips.cc/paper/1989/file/53c3bce66e43be4f209556518c2fcb54-Paper.pdf>`_
-Recurrent Neural Nets (RNN)           deep Learning     ``pyhealth.models.RNN``            2011    `Recurrent neural network based language model <http://www.fit.vutbr.cz/research/groups/speech/servite/2010/rnnlm_mikolov.pdf>`_
-Transformer                           deep Learning     ``pyhealth.models.Transformer``    2017    `Atention is All you Need <https://arxiv.org/abs/1706.03762>`_
-RETAIN                                deep Learning     ``pyhealth.models.RETAIN``         2016    `RETAIN: An Interpretable Predictive Model for Healthcare using Reverse Time Attention Mechanism <https://arxiv.org/abs/1608.05745>`_
-GAMENet                               deep Learning     ``pyhealth.models.GAMENet``        2019    `GAMENet: Graph Attention Mechanism for Explainable Electronic Health Record Prediction <https://arxiv.org/abs/1809.01852>`_
-MICRON                                deep Learning     ``pyhealth.models.MICRON``         2021    `Change Matters: Medication Change Prediction with Recurrent Residual Networks <https://www.ijcai.org/proceedings/2021/0513>`_
-SafeDrug                              deep Learning     ``pyhealth.models.SafeDrug``       2021    `SafeDrug: Dual Molecular Graph Encoders for Recommending Effective and Safe Drug Combinations <https://arxiv.org/abs/2105.02711>`_
-Deepr                                 deep Learning     ``pyhealth.models.Deepr``          2017    `Deepr : A Convolutional Net for Medical Records <https://arxiv.org/abs/1607.07519>`_
-ContraWR Encoder (STFT+CNN)           deep Learning     ``pyhealth.models.ContraWR``       2021    `Self-supervised EEG Representation Learning for Automatic Sleep Staging <https://arxiv.org/abs/2110.15278>`_
-SparcNet (1D CNN)                     deep Learning     ``pyhealth.models.SparcNet``       2023    `Development of Expert-level Classification of Seizures and Rhythmic and Periodic Patterns During EEG Interpretation <#>`_
-==================================    ================  =================================  ======  ============================================================================================================================================================================
+==================================    ================  =================================  ======  ============================================================================================================================================================================  =======================================================================================================================================================================================
+Model Name                            Type              Module                             Year    Summary                                                                                                                                                                       Reference
+==================================    ================  =================================  ======  ============================================================================================================================================================================  =======================================================================================================================================================================================
+Multi-layer Perceptron                deep learning     ``pyhealth.models.MLP``            1986    MLP treats each feature as static                                                                                                                                             `Backpropagation: theory, architectures, and applications <https://www.taylorfrancis.com/books/mono/10.4324/9780203763247/backpropagation-yves-chauvin-david-rumelhart>`_
+Convolutional Neural Network (CNN)    deep learning     ``pyhealth.models.CNN``            1989    CNN runs on the conceptual patient-by-visit grids                                                                                                                             `Handwritten Digit Recognition with a Back-Propagation Network <https://proceedings.neurips.cc/paper/1989/file/53c3bce66e43be4f209556518c2fcb54-Paper.pdf>`_
+Recurrent Neural Nets (RNN)           deep Learning     ``pyhealth.models.RNN``            2011    RNN (includes LSTM and GRU) can run on any sequential level (e.g., visit by visit sequences)                                                                                  `Recurrent neural network based language model <http://www.fit.vutbr.cz/research/groups/speech/servite/2010/rnnlm_mikolov.pdf>`_
+Transformer                           deep Learning     ``pyhealth.models.Transformer``    2017    Transformer can run on any sequential level (e.g., visit by visit sequences)                                                                                                  `Atention is All you Need <https://arxiv.org/abs/1706.03762>`_
+RETAIN                                deep Learning     ``pyhealth.models.RETAIN``         2016    RETAIN uses two RNN to learn patient embeddings while providing feature-level and visit-level importance.                                                                     `RETAIN: An Interpretable Predictive Model for Healthcare using Reverse Time Attention Mechanism <https://arxiv.org/abs/1608.05745>`_
+GAMENet                               deep Learning     ``pyhealth.models.GAMENet``        2019    GAMENet uses memory networks, used only for drug recommendation task                                                                                                          `GAMENet: Graph Attention Mechanism for Explainable Electronic Health Record Prediction <https://arxiv.org/abs/1809.01852>`_
+MICRON                                deep Learning     ``pyhealth.models.MICRON``         2021    MICRON predicts the future drug combination by instead predicting the changes w.r.t. the current combination, used only for drug recommendation task                          `Change Matters: Medication Change Prediction with Recurrent Residual Networks <https://www.ijcai.org/proceedings/2021/0513>`_
+SafeDrug                              deep Learning     ``pyhealth.models.SafeDrug``       2021    SafeDrug encodes drug molecule structures by graph neural networks, used only for drug recommendation task                                                                    `SafeDrug: Dual Molecular Graph Encoders for Recommending Effective and Safe Drug Combinations <https://arxiv.org/abs/2105.02711>`_
+Deepr                                 deep Learning     ``pyhealth.models.Deepr``          2017    Deepr is based on 1D CNN. General purpose.                                                                                                                                    `Deepr : A Convolutional Net for Medical Records <https://arxiv.org/abs/1607.07519>`_
+ContraWR Encoder (STFT+CNN)           deep Learning     ``pyhealth.models.ContraWR``       2021    ContraWR encoder uses short time Fourier transform (STFT) + 2D CNN, used for biosignal learning                                                                               `Self-supervised EEG Representation Learning for Automatic Sleep Staging <https://arxiv.org/abs/2110.15278>`_
+SparcNet (1D CNN)                     deep Learning     ``pyhealth.models.SparcNet``       2023    SparcNet is based on 1D CNN, used for biosignal learning                                                                                                                      `Development of Expert-level Classification of Seizures and Rhythmic and Periodic Patterns During EEG Interpretation <#>`_
+TCN                                   deep learning     ``pyhealth.models.TCN``            2018    TCN is based on dilated 1D CNN. General purpose                                                                                                                               `Temporal Convolutional Networks <https://arxiv.org/abs/1803.01271>`_
+AdaCare                               deep learning     ``pyhealth.models.AdaCare``        2020    AdaCare uses CNNs with dilated filters to learn enriched patient embedding. It uses feature calibration module to provide the feature-level and visit-level interpretability  `AdaCare: Explainable Clinical Health Status Representation Learning via Scale-Adaptive Feature Extraction and Recalibration <https://arxiv.org/abs/1911.12205>`_
+ConCare                               deep learning     ``pyhealth.models.ConCare``        2020    ConCare uses transformers to learn patient embedding and calculate inter-feature correlations.                                                                                `ConCare: Personalized Clinical Feature Embedding via Capturing the Healthcare Context <https://arxiv.org/abs/1911.12216>`_
+StageNet                              deep learning     ``pyhealth.models.StageNet``       2020    StageNet uses stage-aware LSTM to conduct clinical predictive tasks while learning patient disease progression stage change unsupervisedly                                    `StageNet: Stage-Aware Neural Networks for Health Risk Prediction <https://arxiv.org/abs/2001.10054>`_
+Dr. Agent                             deep learning     ``pyhealth.models.Agent``          2020    Dr. Agent uses two reinforcement learning agents to learn patient embeddings by mimicking clinical second opinions                                                            `Dr. Agent: Clinical predictive model via mimicked second opinions <https://academic.oup.com/jamia/article/27/7/1084/5858308>`_
+GRASP                                 deep learning     ``pyhealth.models.GRASP``          2021    GRASP uses graph neural network to identify latent patient clusters and uses the clustering information to learn patient                                                      `GRASP: Generic Framework for Health Status Representation Learning Based on Incorporating Knowledge from Similar Patients <https://ojs.aaai.org/index.php/AAAI/article/view/16152>`_
+==================================    ================  =================================  ======  ============================================================================================================================================================================  =======================================================================================================================================================================================
 
 * Check the `interactive map on benchmark EHR predictive tasks <https://pyhealth.readthedocs.io/en/latest/index.html#benchmark-on-healthcare-tasks>`_.
 
diff --git a/docs/api/models.rst b/docs/api/models.rst
index 9e3328b5..310d0249 100644
--- a/docs/api/models.rst
+++ b/docs/api/models.rst
@@ -18,4 +18,10 @@ We implement the following models for supporting multiple healthcare predictive
     models/pyhealth.models.Deepr
     models/pyhealth.models.ContraWR
     models/pyhealth.models.SparcNet
+    models/pyhealth.models.StageNet
+    models/pyhealth.models.AdaCare
+    models/pyhealth.models.ConCare
+    models/pyhealth.models.Agent
+    models/pyhealth.models.GRASP
+    models/pyhealth.models.TCN
 
diff --git a/docs/api/models/pyhealth.models.AdaCare.rst b/docs/api/models/pyhealth.models.AdaCare.rst
new file mode 100644
index 00000000..00aeaf4f
--- /dev/null
+++ b/docs/api/models/pyhealth.models.AdaCare.rst
@@ -0,0 +1,14 @@
+﻿pyhealth.models.AdaCare
+===================================
+
+The separate callable AdaCareLayer and the complete AdaCare model.
+
+.. autoclass:: pyhealth.models.AdaCareLayer
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.models.AdaCare
+    :members:
+    :undoc-members:
+    :show-inheritance:
\ No newline at end of file
diff --git a/docs/api/models/pyhealth.models.Agent.rst b/docs/api/models/pyhealth.models.Agent.rst
new file mode 100644
index 00000000..e04e282f
--- /dev/null
+++ b/docs/api/models/pyhealth.models.Agent.rst
@@ -0,0 +1,14 @@
+﻿pyhealth.models.Agent
+===================================
+
+The separate callable AgentLayer and the complete Agent model.
+
+.. autoclass:: pyhealth.models.AgentLayer
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.models.Agent
+    :members:
+    :undoc-members:
+    :show-inheritance:
\ No newline at end of file
diff --git a/docs/api/models/pyhealth.models.ConCare.rst b/docs/api/models/pyhealth.models.ConCare.rst
new file mode 100644
index 00000000..d76530bc
--- /dev/null
+++ b/docs/api/models/pyhealth.models.ConCare.rst
@@ -0,0 +1,14 @@
+﻿pyhealth.models.ConCare
+===================================
+
+The separate callable ConCareLayer and the complete ConCare model.
+
+.. autoclass:: pyhealth.models.ConCareLayer
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.models.ConCare
+    :members:
+    :undoc-members:
+    :show-inheritance:
\ No newline at end of file
diff --git a/docs/api/models/pyhealth.models.GRASP.rst b/docs/api/models/pyhealth.models.GRASP.rst
new file mode 100644
index 00000000..7592842e
--- /dev/null
+++ b/docs/api/models/pyhealth.models.GRASP.rst
@@ -0,0 +1,14 @@
+﻿pyhealth.models.GRASP
+===================================
+
+The separate callable GRASPLayer and the complete GRASP model.
+
+.. autoclass:: pyhealth.models.GRASPLayer
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.models.GRASP
+    :members:
+    :undoc-members:
+    :show-inheritance:
\ No newline at end of file
diff --git a/docs/api/models/pyhealth.models.StageNet.rst b/docs/api/models/pyhealth.models.StageNet.rst
new file mode 100644
index 00000000..bf39e9e2
--- /dev/null
+++ b/docs/api/models/pyhealth.models.StageNet.rst
@@ -0,0 +1,14 @@
+﻿pyhealth.models.StageNet
+===================================
+
+The separate callable StageNetLayer and the complete StageNet model.
+
+.. autoclass:: pyhealth.models.StageNetLayer
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.models.StageNet
+    :members:
+    :undoc-members:
+    :show-inheritance:
\ No newline at end of file
diff --git a/docs/api/models/pyhealth.models.TCN.rst b/docs/api/models/pyhealth.models.TCN.rst
new file mode 100644
index 00000000..de1cde81
--- /dev/null
+++ b/docs/api/models/pyhealth.models.TCN.rst
@@ -0,0 +1,14 @@
+﻿pyhealth.models.TCN
+===================================
+
+The separate callable TCNLayer and the complete TCN model.
+
+.. autoclass:: pyhealth.models.TCNLayer
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.models.TCN
+    :members:
+    :undoc-members:
+    :show-inheritance:
\ No newline at end of file
diff --git a/docs/log.rst b/docs/log.rst
index 4c536057..3f651146 100644
--- a/docs/log.rst
+++ b/docs/log.rst
@@ -2,6 +2,11 @@ Development logs
 ======================
 We track the new development here:
 
+**Feb 26, 2023**
+
+.. code-block:: bash
+ 
+    1. add 6 EHR predictiom model: Adacare, Concare, Stagenet, TCN, Grasp, Agent
 
 **Feb 24, 2023**
 
diff --git a/examples/mortality_mimic3_adacare.py b/examples/mortality_mimic3_adacare.py
new file mode 100644
index 00000000..7b5b42fb
--- /dev/null
+++ b/examples/mortality_mimic3_adacare.py
@@ -0,0 +1,51 @@
+from pyhealth.datasets import MIMIC3Dataset
+from pyhealth.datasets import split_by_patient, get_dataloader
+from pyhealth.models import AdaCare
+from pyhealth.tasks import mortality_prediction_mimic3_fn
+from pyhealth.trainer import Trainer
+
+if __name__ == "__main__":
+    # STEP 1: load data
+    base_dataset = MIMIC3Dataset(
+        root="/srv/local/data/physionet.org/files/mimiciii/1.4",
+        tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
+        code_mapping={"ICD9CM": "CCSCM", "ICD9PROC": "CCSPROC", "NDC": "ATC"},
+        dev=False,
+        refresh_cache=False,
+    )
+    base_dataset.stat()
+
+    # STEP 2: set task
+    sample_dataset = base_dataset.set_task(mortality_prediction_mimic3_fn)
+    sample_dataset.stat()
+
+    train_dataset, val_dataset, test_dataset = split_by_patient(
+        sample_dataset, [0.8, 0.1, 0.1]
+    )
+    train_dataloader = get_dataloader(train_dataset, batch_size=32, shuffle=True)
+    val_dataloader = get_dataloader(val_dataset, batch_size=32, shuffle=False)
+    test_dataloader = get_dataloader(test_dataset, batch_size=32, shuffle=False)
+
+    # STEP 3: define model
+    model = AdaCare(
+        dataset=sample_dataset,
+        feature_keys=["conditions", "procedures"],
+        label_key="label",
+        mode="binary",
+        use_embedding=[True, True, True],
+        embedding_dim=32,
+        hidden_dim=32,
+        dropout=0.5,
+    )
+
+    # STEP 4: define trainer
+    trainer = Trainer(model=model)
+    trainer.train(
+        train_dataloader=train_dataloader,
+        val_dataloader=val_dataloader,
+        epochs=50,
+        monitor="roc_auc",
+    )
+
+    # STEP 5: evaluate
+    print(trainer.evaluate(test_dataloader))
diff --git a/examples/mortality_mimic3_agent.py b/examples/mortality_mimic3_agent.py
new file mode 100644
index 00000000..15edd64e
--- /dev/null
+++ b/examples/mortality_mimic3_agent.py
@@ -0,0 +1,49 @@
+from pyhealth.datasets import MIMIC3Dataset
+from pyhealth.datasets import split_by_patient, get_dataloader
+from pyhealth.models import Agent
+from pyhealth.tasks import mortality_prediction_mimic3_fn
+from pyhealth.trainer import Trainer
+
+if __name__ == "__main__":
+    # STEP 1: load data
+    base_dataset = MIMIC3Dataset(
+        root="/srv/local/data/physionet.org/files/mimiciii/1.4",
+        tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
+        code_mapping={"ICD9CM": "CCSCM", "ICD9PROC": "CCSPROC", "NDC": "ATC"},
+        dev=False,
+        refresh_cache=False,
+    )
+    base_dataset.stat()
+
+    # STEP 2: set task
+    sample_dataset = base_dataset.set_task(mortality_prediction_mimic3_fn)
+    sample_dataset.stat()
+
+    train_dataset, val_dataset, test_dataset = split_by_patient(
+        sample_dataset, [0.8, 0.1, 0.1]
+    )
+    train_dataloader = get_dataloader(train_dataset, batch_size=256, shuffle=True)
+    val_dataloader = get_dataloader(val_dataset, batch_size=256, shuffle=False)
+    test_dataloader = get_dataloader(test_dataset, batch_size=256, shuffle=False)
+
+    # STEP 3: define model
+    model = Agent(
+        dataset=sample_dataset,
+        feature_keys=["conditions", "procedures"],
+        label_key="label",
+        mode="binary",
+        embedding_dim=32,
+        hidden_dim=32,
+    )
+
+    # STEP 4: define trainer
+    trainer = Trainer(model=model)
+    trainer.train(
+        train_dataloader=train_dataloader,
+        val_dataloader=val_dataloader,
+        epochs=50,
+        monitor="roc_auc",
+    )
+
+    # STEP 5: evaluate
+    print(trainer.evaluate(test_dataloader))
diff --git a/examples/mortality_mimic3_concare.py b/examples/mortality_mimic3_concare.py
new file mode 100644
index 00000000..58fc6d2d
--- /dev/null
+++ b/examples/mortality_mimic3_concare.py
@@ -0,0 +1,49 @@
+from pyhealth.datasets import MIMIC3Dataset
+from pyhealth.datasets import split_by_patient, get_dataloader
+from pyhealth.models import ConCare
+from pyhealth.tasks import mortality_prediction_mimic3_fn
+from pyhealth.trainer import Trainer
+
+if __name__ == "__main__":
+    # STEP 1: load data
+    base_dataset = MIMIC3Dataset(
+        root="/srv/local/data/physionet.org/files/mimiciii/1.4",
+        tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
+        code_mapping={"ICD9CM": "CCSCM", "ICD9PROC": "CCSPROC", "NDC": "ATC"},
+        dev=False,
+        refresh_cache=False,
+    )
+    base_dataset.stat()
+
+    # STEP 2: set task
+    sample_dataset = base_dataset.set_task(mortality_prediction_mimic3_fn)
+    sample_dataset.stat()
+
+    train_dataset, val_dataset, test_dataset = split_by_patient(
+        sample_dataset, [0.8, 0.1, 0.1]
+    )
+    train_dataloader = get_dataloader(train_dataset, batch_size=256, shuffle=True)
+    val_dataloader = get_dataloader(val_dataset, batch_size=256, shuffle=False)
+    test_dataloader = get_dataloader(test_dataset, batch_size=256, shuffle=False)
+
+    # STEP 3: define model
+    model = ConCare(
+        dataset=sample_dataset,
+        feature_keys=["conditions", "procedures"],
+        label_key="label",
+        mode="binary",
+        use_embedding=[True, True, True],
+        hidden_dim=32,
+    )
+
+    # STEP 4: define trainer
+    trainer = Trainer(model=model)
+    trainer.train(
+        train_dataloader=train_dataloader,
+        val_dataloader=val_dataloader,
+        epochs=3,
+        monitor="roc_auc",
+    )
+
+    # STEP 5: evaluate
+    print(trainer.evaluate(test_dataloader))
diff --git a/examples/mortality_mimic3_grasp.py b/examples/mortality_mimic3_grasp.py
new file mode 100644
index 00000000..011bd13f
--- /dev/null
+++ b/examples/mortality_mimic3_grasp.py
@@ -0,0 +1,50 @@
+from pyhealth.datasets import MIMIC3Dataset
+from pyhealth.datasets import split_by_patient, get_dataloader
+from pyhealth.models import GRASP
+from pyhealth.tasks import mortality_prediction_mimic3_fn
+from pyhealth.trainer import Trainer
+
+if __name__ == "__main__":
+    # STEP 1: load data
+    base_dataset = MIMIC3Dataset(
+        root="/srv/local/data/physionet.org/files/mimiciii/1.4",
+        tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
+        code_mapping={"ICD9CM": "CCSCM", "ICD9PROC": "CCSPROC", "NDC": "ATC"},
+        dev=False,
+        refresh_cache=False,
+    )
+    base_dataset.stat()
+
+    # STEP 2: set task
+    sample_dataset = base_dataset.set_task(mortality_prediction_mimic3_fn)
+    sample_dataset.stat()
+
+    train_dataset, val_dataset, test_dataset = split_by_patient(
+        sample_dataset, [0.8, 0.1, 0.1]
+    )
+    train_dataloader = get_dataloader(train_dataset, batch_size=256, shuffle=True)
+    val_dataloader = get_dataloader(val_dataset, batch_size=256, shuffle=False)
+    test_dataloader = get_dataloader(test_dataset, batch_size=256, shuffle=False)
+
+    # STEP 3: define model
+    model = GRASP(
+        dataset=sample_dataset,
+        feature_keys=["conditions", "procedures"],
+        label_key="label",
+        mode="binary",
+        use_embedding=[True, True, True],
+        embedding_dim=32,
+        hidden_dim=32,
+    )
+
+    # STEP 4: define trainer
+    trainer = Trainer(model=model)
+    trainer.train(
+        train_dataloader=train_dataloader,
+        val_dataloader=val_dataloader,
+        epochs=5,
+        monitor="roc_auc",
+    )
+
+    # STEP 5: evaluate
+    print(trainer.evaluate(test_dataloader))
diff --git a/examples/mortality_mimic3_rnn.py b/examples/mortality_mimic3_rnn.py
index 20cb91a9..5cf7bc9a 100644
--- a/examples/mortality_mimic3_rnn.py
+++ b/examples/mortality_mimic3_rnn.py
@@ -10,7 +10,7 @@
     tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
     code_mapping={"ICD9CM": "CCSCM", "ICD9PROC": "CCSPROC", "NDC": "ATC"},
     dev=False,
-    refresh_cache=True,
+    refresh_cache=False,
 )
 base_dataset.stat()
 
diff --git a/examples/mortality_mimic3_stagenet.py b/examples/mortality_mimic3_stagenet.py
new file mode 100644
index 00000000..52df1186
--- /dev/null
+++ b/examples/mortality_mimic3_stagenet.py
@@ -0,0 +1,48 @@
+from pyhealth.datasets import MIMIC3Dataset
+from pyhealth.datasets import split_by_patient, get_dataloader
+from pyhealth.models import StageNet
+from pyhealth.tasks import mortality_prediction_mimic3_fn
+from pyhealth.trainer import Trainer
+
+if __name__ == "__main__":
+    # STEP 1: load data
+    base_dataset = MIMIC3Dataset(
+        root="/srv/local/data/physionet.org/files/mimiciii/1.4",
+        tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
+        code_mapping={"ICD9CM": "CCSCM", "ICD9PROC": "CCSPROC", "NDC": "ATC"},
+        dev=False,
+        refresh_cache=False,
+    )
+    base_dataset.stat()
+
+    # STEP 2: set task
+    sample_dataset = base_dataset.set_task(mortality_prediction_mimic3_fn)
+    sample_dataset.stat()
+
+    train_dataset, val_dataset, test_dataset = split_by_patient(
+        sample_dataset, [0.8, 0.1, 0.1]
+    )
+    train_dataloader = get_dataloader(train_dataset, batch_size=256, shuffle=True)
+    val_dataloader = get_dataloader(val_dataset, batch_size=256, shuffle=False)
+    test_dataloader = get_dataloader(test_dataset, batch_size=256, shuffle=False)
+
+    # STEP 3: define model
+    model = StageNet(
+        dataset=sample_dataset,
+        feature_keys=["conditions", "procedures"],
+        label_key="label",
+        mode="binary",
+        embedding_dim=32,
+    )
+
+    # STEP 4: define trainer
+    trainer = Trainer(model=model)
+    trainer.train(
+        train_dataloader=train_dataloader,
+        val_dataloader=val_dataloader,
+        epochs=50,
+        monitor="roc_auc",
+    )
+
+    # STEP 5: evaluate
+    print(trainer.evaluate(test_dataloader))
diff --git a/examples/mortality_mimic3_tcn.py b/examples/mortality_mimic3_tcn.py
new file mode 100644
index 00000000..ed1b146a
--- /dev/null
+++ b/examples/mortality_mimic3_tcn.py
@@ -0,0 +1,48 @@
+from pyhealth.datasets import MIMIC3Dataset
+from pyhealth.datasets import split_by_patient, get_dataloader
+from pyhealth.models import TCN
+from pyhealth.tasks import mortality_prediction_mimic3_fn
+from pyhealth.trainer import Trainer
+
+if __name__ == "__main__":
+    # STEP 1: load data
+    base_dataset = MIMIC3Dataset(
+        root="/srv/local/data/physionet.org/files/mimiciii/1.4",
+        tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
+        code_mapping={"ICD9CM": "CCSCM", "ICD9PROC": "CCSPROC", "NDC": "ATC"},
+        dev=False,
+        refresh_cache=False,
+    )
+    base_dataset.stat()
+
+    # STEP 2: set task
+    sample_dataset = base_dataset.set_task(mortality_prediction_mimic3_fn)
+    sample_dataset.stat()
+
+    train_dataset, val_dataset, test_dataset = split_by_patient(
+        sample_dataset, [0.8, 0.1, 0.1]
+    )
+    train_dataloader = get_dataloader(train_dataset, batch_size=256, shuffle=True)
+    val_dataloader = get_dataloader(val_dataset, batch_size=256, shuffle=False)
+    test_dataloader = get_dataloader(test_dataset, batch_size=256, shuffle=False)
+
+    # STEP 3: define model
+    model = TCN(
+        dataset=sample_dataset,
+        feature_keys=["conditions", "procedures"],
+        label_key="label",
+        mode="binary",
+        embedding_dim=32,
+    )
+
+    # STEP 4: define trainer
+    trainer = Trainer(model=model)
+    trainer.train(
+        train_dataloader=train_dataloader,
+        val_dataloader=val_dataloader,
+        epochs=50,
+        monitor="roc_auc",
+    )
+
+    # STEP 5: evaluate
+    print(trainer.evaluate(test_dataloader))
diff --git a/pyhealth/models/__init__.py b/pyhealth/models/__init__.py
index a69f4e6f..af8c8192 100644
--- a/pyhealth/models/__init__.py
+++ b/pyhealth/models/__init__.py
@@ -9,4 +9,16 @@
 from .mlp import MLP
 from .deepr import Deepr, DeeprLayer
 from .contrawr import ResBlock2D, ContraWR
+from .adacare import AdaCare, AdaCareLayer
+from .concare import ConCare, ConCareLayer
+from .agent import Agent, AgentLayer
+from .grasp import GRASP, GRASPLayer
+from .stagenet import StageNet, StageNetLayer
+from .tcn import TCN, TCNLayer
 from .sparcnet import SparcNet, DenseBlock, DenseLayer, TransitionLayer
+from .adacare import AdaCare, AdaCareLayer
+from .concare import ConCare, ConCareLayer
+from .agent import Agent, AgentLayer
+from .grasp import GRASP, GRASPLayer
+from .stagenet import StageNet, StageNetLayer
+from .tcn import TCN, TCNLayer
diff --git a/pyhealth/models/adacare.py b/pyhealth/models/adacare.py
new file mode 100644
index 00000000..1d65c242
--- /dev/null
+++ b/pyhealth/models/adacare.py
@@ -0,0 +1,600 @@
+from typing import List, Tuple, Dict, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.utils.rnn as rnn_utils
+
+from pyhealth.datasets import SampleEHRDataset
+from pyhealth.models import BaseModel
+from pyhealth.models.utils import get_last_visit
+
+
+class Sparsemax(nn.Module):
+    """Sparsemax function."""
+
+    def __init__(self, dim=None):
+        super(Sparsemax, self).__init__()
+
+        self.dim = -1 if dim is None else dim
+
+    def forward(self, input):
+        original_size = input.size()
+        input = input.view(-1, input.size(self.dim))
+
+        dim = 1
+        number_of_logits = input.size(dim)
+
+        input = input - torch.max(input, dim=dim, keepdim=True)[0].expand_as(input)
+
+        zs = torch.sort(input=input, dim=dim, descending=True)[0]
+        range = torch.arange(
+            start=1, end=number_of_logits + 1, dtype=torch.float32
+        ).view(1, -1)
+        range = range.expand_as(zs)
+
+        bound = 1 + range * zs
+        cumulative_sum_zs = torch.cumsum(zs, dim)
+        is_gt = torch.gt(bound, cumulative_sum_zs).type(input.type())
+        k = torch.max(is_gt * range, dim, keepdim=True)[0]
+
+        zs_sparse = is_gt * zs
+        taus = (torch.sum(zs_sparse, dim, keepdim=True) - 1) / k
+        taus = taus.expand_as(input)
+
+        self.output = torch.max(torch.zeros_like(input), input - taus)
+
+        output = self.output.view(original_size)
+
+        return output
+
+    def backward(self, grad_output):
+        dim = 1
+
+        nonzeros = torch.ne(self.output, 0)
+        sum = torch.sum(grad_output * nonzeros, dim=dim) / torch.sum(nonzeros, dim=dim)
+        self.grad_input = nonzeros * (grad_output - sum.expand_as(grad_output))
+
+        return self.grad_input
+
+
+class CausalConv1d(torch.nn.Conv1d):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        dilation=1,
+        groups=1,
+        bias=True,
+    ):
+        self.__padding = (kernel_size - 1) * dilation
+
+        super(CausalConv1d, self).__init__(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=self.__padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+        )
+
+    def forward(self, input):
+        result = super(CausalConv1d, self).forward(input)
+        if self.__padding != 0:
+            return result[:, :, : -self.__padding]
+        return result
+
+
+class Recalibration(nn.Module):
+    def __init__(
+        self, channel, reduction=9, use_h=True, use_c=True, activation="sigmoid"
+    ):
+        super(Recalibration, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool1d(1)
+        self.use_h = use_h
+        self.use_c = use_c
+        scale_dim = 0
+        self.activation = activation
+
+        self.nn_c = nn.Linear(channel, channel // reduction)
+        scale_dim += channel // reduction
+
+        self.nn_rescale = nn.Linear(scale_dim, channel)
+        self.sparsemax = Sparsemax(dim=1)
+        self.softmax = nn.Softmax(dim=1)
+
+    def forward(self, x):
+        b, c, t = x.size()
+
+        y_origin = x.permute(0, 2, 1).reshape(b * t, c).contiguous()
+        se_c = self.nn_c(y_origin)
+        se_c = torch.relu(se_c)
+        y = se_c
+
+        y = self.nn_rescale(y).view(b, t, c).permute(0, 2, 1).contiguous()
+        if self.activation == "sigmoid":
+            y = torch.sigmoid(y)
+        elif self.activation == "sparsemax":
+            y = self.sparsemax(y)
+        else:
+            y = self.softmax(y)
+        return x * y.expand_as(x), y.permute(0, 2, 1)
+
+
+class AdaCareLayer(nn.Module):
+    """AdaCare layer.
+
+    Paper: Liantao Ma et al. Adacare: Explainable clinical health status representation learning
+    via scale-adaptive feature extraction and recalibration. AAAI 2020.
+
+    This layer is used in the AdaCare model. But it can also be used as a
+    standalone layer.
+
+    Args:
+        input_dim: the input feature size.
+        hidden_dim: the hidden dimension of the GRU layer. Default is 128.
+        kernel_size: the kernel size of the causal convolution layer. Default is 2.
+        kernel_num: the kernel number of the causal convolution layer. Default is 64.
+        r_v: the number of the reduction rate for the original feature calibration. Default is 4.
+        r_c: the number of the reduction rate for the convolutional feature recalibration. Default is 4.
+        activation: the activation function for the recalibration layer (sigmoid, sparsemax, softmax). Default is "sigmoid".
+        dropout: dropout rate. Default is 0.5.
+
+    Examples:
+        >>> from pyhealth.models import AdaCareLayer
+        >>> input = torch.randn(3, 128, 64)  # [batch size, sequence len, feature_size]
+        >>> layer = AdaCareLayer(64)
+        >>> c, _, inputatt, convatt = layer(input)
+        >>> c.shape
+        torch.Size([3, 64])
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int = 128,
+        kernel_size: int = 2,
+        kernel_num: int = 64,
+        r_v: int = 4,
+        r_c: int = 4,
+        activation: str = "sigmoid",
+        rnn_type: str = "gru",
+        dropout: float = 0.5,
+    ):
+        super(AdaCareLayer, self).__init__()
+
+        if activation not in ["sigmoid", "softmax", "sparsemax"]:
+            raise ValueError(
+                "Only sigmoid, softmax and sparsemax are supported for activation."
+            )
+        if rnn_type not in ["gru", "lstm"]:
+            raise ValueError("Only gru and lstm are supported for rnn_type.")
+
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.kernel_size = kernel_size
+        self.kernel_num = kernel_num
+        self.r_v = r_v
+        self.r_c = r_c
+        self.dropout = dropout
+
+        self.nn_conv1 = CausalConv1d(input_dim, kernel_num, kernel_size, 1, 1)
+        self.nn_conv3 = CausalConv1d(input_dim, kernel_num, kernel_size, 1, 3)
+        self.nn_conv5 = CausalConv1d(input_dim, kernel_num, kernel_size, 1, 5)
+        torch.nn.init.xavier_uniform_(self.nn_conv1.weight)
+        torch.nn.init.xavier_uniform_(self.nn_conv3.weight)
+        torch.nn.init.xavier_uniform_(self.nn_conv5.weight)
+
+        self.nn_convse = Recalibration(
+            3 * kernel_num, r_c, use_h=False, use_c=True, activation="sigmoid"
+        )
+        self.nn_inputse = Recalibration(
+            input_dim, r_v, use_h=False, use_c=True, activation=activation
+        )
+
+        if rnn_type == "gru":
+            self.rnn = nn.GRU(input_dim + 3 * kernel_num, hidden_dim)
+        else:
+            self.rnn = nn.LSTM(input_dim + 3 * kernel_num, hidden_dim)
+        # self.nn_output = nn.Linear(hidden_dim, output_dim)
+        self.nn_dropout = nn.Dropout(dropout)
+
+        self.relu = nn.ReLU()
+        self.sigmoid = nn.Sigmoid()
+        self.tanh = nn.Tanh()
+
+    def forward(
+        self,
+        x: torch.tensor,
+        mask: Optional[torch.tensor] = None,
+    ) -> Tuple[torch.tensor, torch.tensor, torch.tensor]:
+        """Forward propagation.
+
+        Args:
+            x: a tensor of shape [batch size, sequence len, input_dim].
+            mask: an optional tensor of shape [batch size, sequence len], where
+                1 indicates valid and 0 indicates invalid.
+
+        Returns:
+            last_output: a tensor of shape [batch size, input_dim] representing the
+                patient embedding.
+            output: a tensor of shape [batch size, sequence_len, input_dim] representing the patient embedding at each time step.
+            inputatt: a tensor of shape [batch size, sequence_len, input_dim] representing the feature importance of the input.
+            convatt: a tensor of shape [batch size, sequence_len, 3 * kernel_num] representing the feature importance of the convolutional features.
+        """
+        # batch_size = x.size(0)
+        # time_step = x.size(1)
+        # feature_dim = x.size(2)
+
+        conv_input = x.permute(0, 2, 1)
+        conv_res1 = self.nn_conv1(conv_input)
+        conv_res3 = self.nn_conv3(conv_input)
+        conv_res5 = self.nn_conv5(conv_input)
+
+        conv_res = torch.cat((conv_res1, conv_res3, conv_res5), dim=1)
+        conv_res = self.relu(conv_res)
+
+        convse_res, convatt = self.nn_convse(conv_res)
+        inputse_res, inputatt = self.nn_inputse(x.permute(0, 2, 1))
+        concat_input = torch.cat((convse_res, inputse_res), dim=1).permute(0, 2, 1)
+        output, _ = self.rnn(concat_input)
+        last_output = get_last_visit(output, mask)
+        if self.dropout > 0.0:
+            last_output = self.nn_dropout(last_output)
+        return last_output, output, inputatt, convatt
+
+
+class AdaCare(BaseModel):
+    """AdaCare model.
+
+    Paper: Liantao Ma et al. Adacare: Explainable clinical health status representation learning
+    via scale-adaptive feature extraction and recalibration. AAAI 2020.
+
+    Note:
+        We use separate AdaCare layers for different feature_keys.
+        Currently, we automatically support different input formats:
+            - code based input (need to use the embedding table later)
+            - float/int based value input
+        Since the AdaCare model calibrate the original features to provide interpretability, we do not recommend use embeddings for the input features.
+        We follow the current convention for the AdaCare model:
+            - case 1. [code1, code2, code3, ...]
+                - we will assume the code follows the order; our model will encode
+                each code into a vector and apply AdaCare on the code level
+            - case 2. [[code1, code2]] or [[code1, code2], [code3, code4, code5], ...]
+                - we will assume the inner bracket follows the order; our model first
+                use the embedding table to encode each code into a vector and then use
+                average/mean pooling to get one vector for one inner bracket; then use
+                AdaCare one the braket level
+            - case 3. [[1.5, 2.0, 0.0]] or [[1.5, 2.0, 0.0], [8, 1.2, 4.5], ...]
+                - this case only makes sense when each inner bracket has the same length;
+                we assume each dimension has the same meaning; we run AdaCare directly
+                on the inner bracket level, similar to case 1 after embedding table
+            - case 4. [[[1.5, 2.0, 0.0]]] or [[[1.5, 2.0, 0.0], [8, 1.2, 4.5]], ...]
+                - this case only makes sense when each inner bracket has the same length;
+                we assume each dimension has the same meaning; we run AdaCare directly
+                on the inner bracket level, similar to case 2 after embedding table
+
+    Args:
+        dataset: the dataset to train the model. It is used to query certain
+            information such as the set of all tokens.
+        feature_keys:  list of keys in samples to use as features,
+            e.g. ["conditions", "procedures"].
+        label_key: key in samples to use as label (e.g., "drugs").
+        mode: one of "binary", "multiclass", or "multilabel".
+        use_embedding: list of bools indicating whether to use embedding for each feature type,
+            e.g. [True, False].
+        embedding_dim: the embedding dimension. Default is 128.
+        hidden_dim: the hidden dimension. Default is 128.
+        **kwargs: other parameters for the AdaCare layer.
+
+
+    Examples:
+        >>> from pyhealth.datasets import SampleEHRDataset
+        >>> samples = [
+        ...         {
+        ...             "patient_id": "patient-0",
+        ...             "visit_id": "visit-0",
+        ...             "list_codes": ["505800458", "50580045810", "50580045811"],  # NDC
+        ...             "list_vectors": [[1.0, 2.55, 3.4], [4.1, 5.5, 6.0]],
+        ...             "list_list_codes": [["A05B", "A05C", "A06A"], ["A11D", "A11E"]],  # ATC-4
+        ...             "list_list_vectors": [
+        ...                 [[1.8, 2.25, 3.41], [4.50, 5.9, 6.0]],
+        ...                 [[7.7, 8.5, 9.4]],
+        ...             ],
+        ...             "label": 1,
+        ...         },
+        ...         {
+        ...             "patient_id": "patient-0",
+        ...             "visit_id": "visit-1",
+        ...             "list_codes": [
+        ...                 "55154191800",
+        ...                 "551541928",
+        ...                 "55154192800",
+        ...                 "705182798",
+        ...                 "70518279800",
+        ...             ],
+        ...             "list_vectors": [[1.4, 3.2, 3.5], [4.1, 5.9, 1.7], [4.5, 5.9, 1.7]],
+        ...             "list_list_codes": [["A04A", "B035", "C129"]],
+        ...             "list_list_vectors": [
+        ...                 [[1.0, 2.8, 3.3], [4.9, 5.0, 6.6], [7.7, 8.4, 1.3], [7.7, 8.4, 1.3]],
+        ...             ],
+        ...             "label": 0,
+        ...         },
+        ...     ]
+        >>> dataset = SampleEHRDataset(samples=samples, dataset_name="test")
+        >>>
+        >>> from pyhealth.models import AdaCare
+        >>> model = AdaCare(
+        ...         dataset=dataset,
+        ...         feature_keys=[
+        ...             "list_codes",
+        ...             "list_vectors",
+        ...             "list_list_codes",
+        ...             "list_list_vectors",
+        ...         ],
+        ...         label_key="label",
+        ...         use_embedding=[True, False, True, False],
+        ...         mode="binary",
+        ...     )
+        >>>
+        >>> from pyhealth.datasets import get_dataloader
+        >>> train_loader = get_dataloader(dataset, batch_size=2, shuffle=True)
+        >>> data_batch = next(iter(train_loader))
+        >>>
+        >>> ret = model(**data_batch)
+        >>> print(ret)
+        {'loss': tensor(0.7234, grad_fn=<BinaryCrossEntropyWithLogitsBackward0>), 'feature_importance': tensor(...), 'conv_feature_importance': tensor(...), 'y_prob': tensor([[0.5423],
+                [0.5142]], grad_fn=<SigmoidBackward0>), 'y_true': tensor([[0.],
+                [1.]])}
+        >>>
+
+    """
+
+    def __init__(
+        self,
+        dataset: SampleEHRDataset,
+        feature_keys: List[str],
+        label_key: str,
+        mode: str,
+        use_embedding: List[bool],
+        embedding_dim: int = 128,
+        hidden_dim: int = 128,
+        **kwargs,
+    ):
+        super(AdaCare, self).__init__(
+            dataset=dataset,
+            feature_keys=feature_keys,
+            label_key=label_key,
+            mode=mode,
+        )
+        self.embedding_dim = embedding_dim
+        self.use_embedding = use_embedding
+        self.hidden_dim = hidden_dim
+
+        if "input_dim" in kwargs:
+            raise ValueError("input_dim is automatically determined")
+
+        # the key of self.feat_tokenizers only contains the code based inputs
+        self.feat_tokenizers = {}
+        self.label_tokenizer = self.get_label_tokenizer()
+        # the key of self.embeddings only contains the code based inputs
+        self.embeddings = nn.ModuleDict()
+        # the key of self.linear_layers only contains the float/int based inputs
+        self.linear_layers = nn.ModuleDict()
+        self.adacare = nn.ModuleDict()
+
+        for idx, feature_key in enumerate(self.feature_keys):
+            input_info = self.dataset.input_info[feature_key]
+            # sanity check
+            if input_info["type"] not in [str, float, int]:
+                raise ValueError(
+                    "AdaCare only supports str code, float and int as input types"
+                )
+            elif (input_info["type"] == str) and (input_info["dim"] not in [2, 3]):
+                raise ValueError(
+                    "AdaCare only supports 2-dim or 3-dim str code as input types"
+                )
+            elif (input_info["type"] == str) and (use_embedding[idx] == False):
+                raise ValueError(
+                    "AdaCare only supports embedding for str code as input types"
+                )
+            elif (input_info["type"] in [float, int]) and (
+                input_info["dim"] not in [2, 3]
+            ):
+                raise ValueError(
+                    "AdaCare only supports 2-dim or 3-dim float and int as input types"
+                )
+
+            # for code based input, we need Type
+            # for float/int based input, we need Type, input_dim
+            if use_embedding[idx]:
+                self.add_feature_transform_layer(feature_key, input_info)
+                self.adacare[feature_key] = AdaCareLayer(
+                    input_dim=embedding_dim, hidden_dim=self.hidden_dim, **kwargs
+                )
+            else:
+                self.adacare[feature_key] = AdaCareLayer(
+                    input_dim=input_info["len"], hidden_dim=self.hidden_dim, **kwargs
+                )
+
+        output_size = self.get_output_size(self.label_tokenizer)
+        self.fc = nn.Linear(len(self.feature_keys) * self.hidden_dim, output_size)
+
+    def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
+        """Forward propagation.
+
+        The label `kwargs[self.label_key]` is a list of labels for each patient.
+
+        Args:
+            **kwargs: keyword arguments for the model. The keys must contain
+                all the feature keys and the label key.
+
+        Returns:
+            A dictionary with the following keys:
+                loss: a scalar tensor representing the loss.
+                feature_importance: a list of tensors with shape (feature_type, batch_size, time_step, features)
+                                    representing the feature importance.
+                conv_feature_importance: a list of tensors with shape (feature_type, batch_size, time_step, 3*kernal_size)
+                                        representing the convolutional feature importance.
+                y_prob: a tensor representing the predicted probabilities.
+                y_true: a tensor representing the true labels.
+        """
+        patient_emb = []
+        feature_importance = []
+        conv_feature_importance = []
+        for idx, feature_key in enumerate(self.feature_keys):
+            input_info = self.dataset.input_info[feature_key]
+            dim_, type_ = input_info["dim"], input_info["type"]
+
+            # for case 1: [code1, code2, code3, ...]
+            if (dim_ == 2) and (type_ == str):
+                x = self.feat_tokenizers[feature_key].batch_encode_2d(
+                    kwargs[feature_key]
+                )
+                # (patient, event)
+                x = torch.tensor(x, dtype=torch.long, device=self.device)
+                # (patient, event, embedding_dim)
+                x = self.embeddings[feature_key](x)
+                # (patient, event)
+                mask = torch.sum(x, dim=2) != 0
+
+            # for case 2: [[code1, code2], [code3, ...], ...]
+            elif (dim_ == 3) and (type_ == str):
+                x = self.feat_tokenizers[feature_key].batch_encode_3d(
+                    kwargs[feature_key]
+                )
+                # (patient, visit, event)
+                x = torch.tensor(x, dtype=torch.long, device=self.device)
+                # (patient, visit, event, embedding_dim)
+                x = self.embeddings[feature_key](x)
+                # (patient, visit, embedding_dim)
+                x = torch.sum(x, dim=2)
+                # (patient, visit)
+                mask = torch.sum(x, dim=2) != 0
+
+            # for case 3: [[1.5, 2.0, 0.0], ...]
+            elif (dim_ == 2) and (type_ in [float, int]):
+                x, mask = self.padding2d(kwargs[feature_key])
+                # (patient, event, values)
+                x = torch.tensor(x, dtype=torch.float, device=self.device)
+                # (patient, event, embedding_dim)
+                if self.use_embedding[idx]:
+                    x = self.linear_layers[feature_key](x)
+                # (patient, event)
+                mask = mask.bool().to(self.device)
+
+            # for case 4: [[[1.5, 2.0, 0.0], [1.8, 2.4, 6.0]], ...]
+            elif (dim_ == 3) and (type_ in [float, int]):
+                x, mask = self.padding3d(kwargs[feature_key])
+                # (patient, visit, event, values)
+                x = torch.tensor(x, dtype=torch.float, device=self.device)
+                # (patient, visit, embedding_dim)
+                x = torch.sum(x, dim=2)
+
+                if self.use_embedding[idx]:
+                    x = self.linear_layers[feature_key](x)
+                # (patient, event)
+                mask = mask[:, :, 0]
+                mask = mask.bool().to(self.device)
+
+            else:
+                raise NotImplementedError
+
+            x, _, inputatt, convatt = self.adacare[feature_key](x, mask)
+            feature_importance.append(inputatt)
+            conv_feature_importance.append(convatt)
+            patient_emb.append(x)
+
+        patient_emb = torch.cat(patient_emb, dim=1)
+        # (patient, label_size)
+        logits = self.fc(patient_emb)
+        # obtain y_true, loss, y_prob
+        y_true = self.prepare_labels(kwargs[self.label_key], self.label_tokenizer)
+        loss = self.get_loss_function()(logits, y_true)
+        y_prob = self.prepare_y_prob(logits)
+        # return {
+        #     "loss": loss,
+        #     "feature_importance": feature_importance,
+        #     "conv_feature_importance": conv_feature_importance,
+        #     "y_prob": y_prob,
+        #     "y_true": y_true,
+        # }
+        return {
+            "loss": loss,
+            "y_prob": y_prob,
+            "y_true": y_true,
+        }
+
+
+if __name__ == "__main__":
+    from pyhealth.datasets import SampleEHRDataset
+
+    samples = [
+        {
+            "patient_id": "patient-0",
+            "visit_id": "visit-0",
+            # "single_vector": [1, 2, 3],
+            "list_codes": ["505800458", "50580045810", "50580045811"],  # NDC
+            "list_vectors": [[1.0, 2.55, 3.4], [4.1, 5.5, 6.0]],
+            "list_list_codes": [["A05B", "A05C", "A06A"], ["A11D", "A11E"]],  # ATC-4
+            "list_list_vectors": [
+                [[1.8, 2.25, 3.41], [4.50, 5.9, 6.0]],
+                [[7.7, 8.5, 9.4]],
+            ],
+            "label": 1,
+        },
+        {
+            "patient_id": "patient-0",
+            "visit_id": "visit-1",
+            # "single_vector": [1, 5, 8],
+            "list_codes": [
+                "55154191800",
+                "551541928",
+                "55154192800",
+                "705182798",
+                "70518279800",
+            ],
+            "list_vectors": [[1.4, 3.2, 3.5], [4.1, 5.9, 1.7], [4.5, 5.9, 1.7]],
+            "list_list_codes": [["A04A", "B035", "C129"]],
+            "list_list_vectors": [
+                [[1.0, 2.8, 3.3], [4.9, 5.0, 6.6], [7.7, 8.4, 1.3], [7.7, 8.4, 1.3]],
+            ],
+            "label": 0,
+        },
+    ]
+
+    # dataset
+    dataset = SampleEHRDataset(samples=samples, dataset_name="test")
+
+    # data loader
+    from pyhealth.datasets import get_dataloader
+
+    train_loader = get_dataloader(dataset, batch_size=2, shuffle=True)
+
+    # model
+    model = AdaCare(
+        dataset=dataset,
+        feature_keys=[
+            "list_codes",
+            "list_vectors",
+            "list_list_codes",
+            "list_list_vectors",
+        ],
+        label_key="label",
+        mode="binary",
+        use_embedding=[True, False, True, False],
+        hidden_dim=64,
+    )
+
+    # data batch
+    data_batch = next(iter(train_loader))
+
+    # try the model
+    ret = model(**data_batch)
+    print(ret)
+
+    # try loss backward
+    ret["loss"].backward()
diff --git a/pyhealth/models/agent.py b/pyhealth/models/agent.py
new file mode 100644
index 00000000..fee11ac0
--- /dev/null
+++ b/pyhealth/models/agent.py
@@ -0,0 +1,739 @@
+from typing import List, Tuple, Dict, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.utils.rnn as rnn_utils
+
+from pyhealth.datasets import SampleEHRDataset
+from pyhealth.models import BaseModel
+from pyhealth.models.utils import get_last_visit
+import torch.nn.functional as F
+
+
+class AgentLayer(nn.Module):
+    """Dr. Agent layer.
+
+    Paper: Junyi Gao et al. Dr. Agent: Clinical predictive model via mimicked second opinions. JAMIA.
+
+    This layer is used in the Dr. Agent model. But it can also be used as a
+    standalone layer.
+
+    Args:
+        input_dim: dynamic feature size.
+        static_dim: static feature size, if 0, then no static feature is used.
+        cell: rnn cell type. Default is "gru".
+        use_baseline: whether to use baseline for the RL agent. Default is True.
+        n_actions: number of historical visits to choose. Default is 10.
+        n_units: number of hidden units in each agent. Default is 64.
+        fusion_dim: number of hidden units in the final representation. Default is 128.
+        n_hidden: number of hidden units in the rnn. Default is 128.
+        dropout: dropout rate. Default is 0.5.
+        lamda: weight for the agent selected hidden state and the current hidden state. Default is 0.5.
+
+    Examples:
+        >>> from pyhealth.models import AgentLayer
+        >>> input = torch.randn(3, 128, 64)  # [batch size, sequence len, feature_size]
+        >>> layer = AgentLayer(64)
+        >>> c, _ = layer(input)
+        >>> c.shape
+        torch.Size([3, 128])
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        static_dim: int = 0,
+        cell: str = "gru",
+        use_baseline: bool = True,
+        n_actions: int = 10,
+        n_units: int = 64,
+        n_hidden: int = 128,
+        dropout: int = 0.5,
+        lamda: int = 0.5,
+    ):
+        super(AgentLayer, self).__init__()
+
+        if cell not in ["gru", "lstm"]:
+            raise ValueError("Only gru and lstm are supported for cell.")
+
+        self.cell = cell
+        self.use_baseline = use_baseline
+        self.n_actions = n_actions
+        self.n_units = n_units
+        self.input_dim = input_dim
+        self.n_hidden = n_hidden
+        # self.n_output = n_output
+        self.dropout = dropout
+        self.lamda = lamda
+        self.fusion_dim = n_hidden
+        self.static_dim = static_dim
+
+        self.agent1_action = []
+        self.agent1_prob = []
+        self.agent1_entropy = []
+        self.agent1_baseline = []
+        self.agent2_action = []
+        self.agent2_prob = []
+        self.agent2_entropy = []
+        self.agent2_baseline = []
+
+        self.agent1_fc1 = nn.Linear(self.n_hidden + self.static_dim, self.n_units)
+        self.agent2_fc1 = nn.Linear(self.input_dim + self.static_dim, self.n_units)
+        self.agent1_fc2 = nn.Linear(self.n_units, self.n_actions)
+        self.agent2_fc2 = nn.Linear(self.n_units, self.n_actions)
+        if use_baseline == True:
+            self.agent1_value = nn.Linear(self.n_units, 1)
+            self.agent2_value = nn.Linear(self.n_units, 1)
+
+        if self.cell == "lstm":
+            self.rnn = nn.LSTMCell(self.input_dim, self.n_hidden)
+        else:
+            self.rnn = nn.GRUCell(self.input_dim, self.n_hidden)
+
+        for name, param in self.rnn.named_parameters():
+            if "bias" in name:
+                nn.init.constant_(param, 0.0)
+            elif "weight" in name:
+                nn.init.orthogonal_(param)
+
+        if dropout > 0.0:
+            self.nn_dropout = nn.Dropout(p=dropout)
+        if self.static_dim > 0:
+            self.init_h = nn.Linear(self.static_dim, self.n_hidden)
+            self.init_c = nn.Linear(self.static_dim, self.n_hidden)
+            self.fusion = nn.Linear(self.n_hidden + self.static_dim, self.fusion_dim)
+        # self.output = nn.Linear(self.fusion_dim, self.n_output)
+
+        self.sigmoid = nn.Sigmoid()
+        self.softmax = nn.Softmax(dim=1)
+        self.tanh = nn.Tanh()
+        self.relu = nn.ReLU()
+
+    def choose_action(self, observation, agent=1):
+        observation = observation.detach()
+
+        if agent == 1:
+            result_fc1 = self.agent1_fc1(observation)
+            result_fc1 = self.tanh(result_fc1)
+            result_fc2 = self.agent1_fc2(result_fc1)
+            if self.use_baseline == True:
+                result_value = self.agent1_value(result_fc1)
+                self.agent1_baseline.append(result_value)
+        else:
+            result_fc1 = self.agent2_fc1(observation)
+            result_fc1 = self.tanh(result_fc1)
+            result_fc2 = self.agent2_fc2(result_fc1)
+            if self.use_baseline == True:
+                result_value = self.agent2_value(result_fc1)
+                self.agent2_baseline.append(result_value)
+
+        probs = self.softmax(result_fc2)
+        m = torch.distributions.Categorical(probs)
+        actions = m.sample()
+
+        if agent == 1:
+            self.agent1_entropy.append(m.entropy())
+            self.agent1_action.append(actions.unsqueeze(-1))
+            self.agent1_prob.append(m.log_prob(actions))
+        else:
+            self.agent2_entropy.append(m.entropy())
+            self.agent2_action.append(actions.unsqueeze(-1))
+            self.agent2_prob.append(m.log_prob(actions))
+
+        return actions.unsqueeze(-1)
+
+    def forward(
+        self,
+        x: torch.tensor,
+        static: Optional[torch.tensor] = None,
+        mask: Optional[torch.tensor] = None,
+    ) -> Tuple[torch.tensor]:
+        """Forward propagation.
+
+        Args:
+            x: a tensor of shape [batch size, sequence len, input_dim].
+            static: a tensor of shape [batch size, static_dim].
+            mask: an optional tensor of shape [batch size, sequence len], where
+                1 indicates valid and 0 indicates invalid.
+
+        Returns:
+            last_output: a tensor of shape [batch size, n_hidden] representing the
+                patient embedding.
+            output: a tensor of shape [batch size, sequence len, n_hidden] representing the patient embedding at each time step.
+        """
+        # rnn will only apply dropout between layers
+
+        batch_size = x.size(0)
+        time_step = x.size(1)
+        feature_dim = x.size(2)
+
+        self.agent1_action = []
+        self.agent1_prob = []
+        self.agent1_entropy = []
+        self.agent1_baseline = []
+        self.agent2_action = []
+        self.agent2_prob = []
+        self.agent2_entropy = []
+        self.agent2_baseline = []
+
+        if self.static_dim > 0:
+            cur_h = self.init_h(static)
+            if self.cell == "lstm":
+                cur_c = self.init_c(static)
+        else:
+            cur_h = torch.zeros(
+                batch_size, self.n_hidden, dtype=torch.float32, device=x.device
+            )
+            if self.cell == "lstm":
+                cur_c = torch.zeros(
+                    batch_size, self.n_hidden, dtype=torch.float32, device=x.device
+                )
+
+        h = []
+        for cur_time in range(time_step):
+            cur_input = x[:, cur_time, :]
+
+            if cur_time == 0:
+                obs_1 = cur_h
+                obs_2 = cur_input
+
+                if self.static_dim > 0:
+                    obs_1 = torch.cat((obs_1, static), dim=1)
+                    obs_2 = torch.cat((obs_2, static), dim=1)
+
+                self.choose_action(obs_1, 1).long()
+                self.choose_action(obs_2, 2).long()
+
+                observed_h = (
+                    torch.zeros_like(cur_h, dtype=torch.float32)
+                    .view(-1)
+                    .repeat(self.n_actions)
+                    .view(self.n_actions, batch_size, self.n_hidden)
+                )
+                action_h = cur_h
+                if self.cell == "lstm":
+                    observed_c = (
+                        torch.zeros_like(cur_c, dtype=torch.float32)
+                        .view(-1)
+                        .repeat(self.n_actions)
+                        .view(self.n_actions, batch_size, self.n_hidden)
+                    )
+                    action_c = cur_c
+
+            else:
+                observed_h = torch.cat((observed_h[1:], cur_h.unsqueeze(0)), 0)
+
+                obs_1 = observed_h.mean(dim=0)
+                obs_2 = cur_input
+
+                if self.static_dim > 0:
+                    obs_1 = torch.cat((obs_1, static), dim=1)
+                    obs_2 = torch.cat((obs_2, static), dim=1)
+
+                act_idx1 = self.choose_action(obs_1, 1).long()
+                act_idx2 = self.choose_action(obs_2, 2).long()
+                batch_idx = torch.arange(batch_size, dtype=torch.long).unsqueeze(-1)
+                action_h1 = observed_h[act_idx1, batch_idx, :].squeeze(1)
+                action_h2 = observed_h[act_idx2, batch_idx, :].squeeze(1)
+                action_h = (action_h1 + action_h2) / 2
+                if self.cell == "lstm":
+                    observed_c = torch.cat((observed_c[1:], cur_c.unsqueeze(0)), 0)
+                    action_c1 = observed_c[act_idx1, batch_idx, :].squeeze(1)
+                    action_c2 = observed_c[act_idx2, batch_idx, :].squeeze(1)
+                    action_c = (action_c1 + action_c2) / 2
+
+            if self.cell == "lstm":
+                weighted_h = self.lamda * action_h + (1 - self.lamda) * cur_h
+                weighted_c = self.lamda * action_c + (1 - self.lamda) * cur_c
+                rnn_state = (weighted_h, weighted_c)
+                cur_h, cur_c = self.rnn(cur_input, rnn_state)
+            else:
+                weighted_h = self.lamda * action_h + (1 - self.lamda) * cur_h
+                cur_h = self.rnn(cur_input, weighted_h)
+            h.append(cur_h)
+
+        h = torch.stack(h, dim=1)
+
+        if self.static_dim > 0:
+            static = static.unsqueeze(1).repeat(1, time_step, 1)
+            h = torch.cat((h, static), dim=2)
+            h = self.fusion(h)
+
+        last_out = get_last_visit(h, mask)
+
+        if self.dropout > 0.0:
+            last_out = self.nn_dropout(last_out)
+        return last_out, h
+
+
+class Agent(BaseModel):
+    """Dr. Agent model.
+
+    Paper: Junyi Gao et al. Dr. Agent: Clinical predictive model via mimicked second opinions. JAMIA.
+
+    Note:
+        We use separate Dr. Agent layers for different feature_keys.
+        Currently, we automatically support different input formats:
+            - code based input (need to use the embedding table later)
+            - float/int based value input
+        We follow the current convention for the Dr. Agent model:
+            - case 1. [code1, code2, code3, ...]
+                - we will assume the code follows the order; our model will encode
+                each code into a vector and apply Dr. Agent on the code level
+            - case 2. [[code1, code2]] or [[code1, code2], [code3, code4, code5], ...]
+                - we will assume the inner bracket follows the order; our model first
+                use the embedding table to encode each code into a vector and then use
+                average/mean pooling to get one vector for one inner bracket; then use
+                Dr. Agent one the braket level
+            - case 3. [[1.5, 2.0, 0.0]] or [[1.5, 2.0, 0.0], [8, 1.2, 4.5], ...]
+                - this case only makes sense when each inner bracket has the same length;
+                we assume each dimension has the same meaning; we run Dr. Agent directly
+                on the inner bracket level, similar to case 1 after embedding table
+            - case 4. [[[1.5, 2.0, 0.0]]] or [[[1.5, 2.0, 0.0], [8, 1.2, 4.5]], ...]
+                - this case only makes sense when each inner bracket has the same length;
+                we assume each dimension has the same meaning; we run Dr. Agent directly
+                on the inner bracket level, similar to case 2 after embedding table
+
+    Args:
+        dataset: the dataset to train the model. It is used to query certain
+            information such as the set of all tokens.
+        feature_keys:  list of keys in samples to use as features,
+            e.g. ["conditions", "procedures"].
+        label_key: key in samples to use as label (e.g., "drugs").
+        mode: one of "binary", "multiclass", or "multilabel".
+        static_keys: the key in samples to use as static features, e.g. "demographics". Default is None.
+                     we only support numerical static features.
+        embedding_dim: the embedding dimension. Default is 128.
+        hidden_dim: the hidden dimension of the RNN in the Dr. Agent layer. Default is 128.
+        use_baseline: whether to use the baseline value to calculate the RL loss. Default is True.
+        **kwargs: other parameters for the Dr. Agent layer.
+
+
+    Examples:
+        >>> from pyhealth.datasets import SampleEHRDataset
+        >>> samples = [
+        ...         {
+        ...             "patient_id": "patient-0",
+        ...             "visit_id": "visit-0",
+        ...             "list_codes": ["505800458", "50580045810", "50580045811"],  # NDC
+        ...             "list_vectors": [[1.0, 2.55, 3.4], [4.1, 5.5, 6.0]],
+        ...             "list_list_codes": [["A05B", "A05C", "A06A"], ["A11D", "A11E"]],  # ATC-4
+        ...             "list_list_vectors": [
+        ...                 [[1.8, 2.25, 3.41], [4.50, 5.9, 6.0]],
+        ...                 [[7.7, 8.5, 9.4]],
+        ...             ],
+        ...             "demographic": [0.0, 2.0, 1.5],
+        ...             "label": 1,
+        ...         },
+        ...         {
+        ...             "patient_id": "patient-0",
+        ...             "visit_id": "visit-1",
+        ...             "list_codes": [
+        ...                 "55154191800",
+        ...                 "551541928",
+        ...                 "55154192800",
+        ...                 "705182798",
+        ...                 "70518279800",
+        ...             ],
+        ...             "list_vectors": [[1.4, 3.2, 3.5], [4.1, 5.9, 1.7], [4.5, 5.9, 1.7]],
+        ...             "list_list_codes": [["A04A", "B035", "C129"]],
+        ...             "list_list_vectors": [
+        ...                 [[1.0, 2.8, 3.3], [4.9, 5.0, 6.6], [7.7, 8.4, 1.3], [7.7, 8.4, 1.3]],
+        ...             ],
+        ...             "demographic": [0.0, 2.0, 1.5],
+        ...             "label": 0,
+        ...         },
+        ...     ]
+        >>> dataset = SampleEHRDataset(samples=samples, dataset_name="test")
+        >>>
+        >>> from pyhealth.models import Agent
+        >>> model = Agent(
+        ...         dataset=dataset,
+        ...         feature_keys=[
+        ...             "list_codes",
+        ...             "list_vectors",
+        ...             "list_list_codes",
+        ...             "list_list_vectors",
+        ...         ],
+        ...         label_key="label",
+        ...         static_key="demographic",
+        ...         mode="binary"
+        ...     )
+        >>>
+        >>> from pyhealth.datasets import get_dataloader
+        >>> train_loader = get_dataloader(dataset, batch_size=2, shuffle=True)
+        >>> data_batch = next(iter(train_loader))
+        >>>
+        >>> ret = model(**data_batch)
+        >>> print(ret)
+        {'loss': tensor(8.0580, grad_fn=<AddBackward0>), 'loss_task': tensor(0.6588, grad_fn=<BinaryCrossEntropyWithLogitsBackward0>), 'loss_rl': tensor(7.3992, grad_fn=<AddBackward0>), 'y_prob': tensor([[0.4761], [0.4376]], grad_fn=<SigmoidBackward0>), 'y_true': tensor([[1.], [0.]])}
+        >>>
+
+    """
+
+    def __init__(
+        self,
+        dataset: SampleEHRDataset,
+        feature_keys: List[str],
+        label_key: str,
+        mode: str,
+        static_key: Optional[str] = None,
+        embedding_dim: int = 128,
+        hidden_dim: int = 128,
+        use_baseline: bool = True,
+        **kwargs,
+    ):
+        super(Agent, self).__init__(
+            dataset=dataset,
+            feature_keys=feature_keys,
+            label_key=label_key,
+            mode=mode,
+        )
+        self.embedding_dim = embedding_dim
+        self.hidden_dim = hidden_dim
+
+        # validate kwargs for Dr. Agent layer
+        if "feature_size" in kwargs:
+            raise ValueError("feature_size is determined by embedding_dim")
+
+        # the key of self.feat_tokenizers only contains the code based inputs
+        self.feat_tokenizers = {}
+        self.static_key = static_key
+        self.use_baseline = use_baseline
+        self.label_tokenizer = self.get_label_tokenizer()
+        # the key of self.embeddings only contains the code based inputs
+        self.embeddings = nn.ModuleDict()
+        # the key of self.linear_layers only contains the float/int based inputs
+        self.linear_layers = nn.ModuleDict()
+
+        self.static_dim = 0
+        if self.static_key is not None:
+            self.static_dim = self.dataset.input_info[self.static_key]["len"]
+
+        self.agent = nn.ModuleDict()
+        # add feature Dr. Agent layers
+        for feature_key in self.feature_keys:
+            input_info = self.dataset.input_info[feature_key]
+            # sanity check
+            if input_info["type"] not in [str, float, int]:
+                raise ValueError(
+                    "Dr. Agent only supports str code, float and int as input types"
+                )
+            elif (input_info["type"] == str) and (input_info["dim"] not in [2, 3]):
+                raise ValueError(
+                    "Dr. Agent only supports 2-dim or 3-dim str code as input types"
+                )
+            elif (input_info["type"] in [float, int]) and (
+                input_info["dim"] not in [2, 3]
+            ):
+                raise ValueError(
+                    "Dr. Agent only supports 2-dim or 3-dim float and int as input types"
+                )
+
+            # for code based input, we need Type
+            # for float/int based input, we need Type, input_dim
+            self.add_feature_transform_layer(feature_key, input_info)
+            self.agent[feature_key] = AgentLayer(
+                input_dim=embedding_dim,
+                static_dim=self.static_dim,
+                n_hidden=hidden_dim,
+                **kwargs,
+            )
+
+        output_size = self.get_output_size(self.label_tokenizer)
+        self.fc = nn.Linear(len(self.feature_keys) * self.hidden_dim, output_size)
+
+    def get_loss(self, model, pred, true, mask, gamma=0.9, entropy_term=0.01):
+
+        if self.mode == "binary":
+            pred = torch.sigmoid(pred)
+            rewards = ((pred - 0.5) * 2 * true).squeeze()
+        elif self.mode == "multiclass":
+            pred = torch.softmax(pred, dim=-1)
+            y_onehot = torch.zeros_like(pred).scatter(1, true.unsqueeze(1), 1)
+            rewards = (pred * y_onehot).sum(-1).squeeze()
+        elif self.mode == "multilabel":
+            pred = torch.sigmoid(pred)
+            rewards = (
+                ((pred - 0.5) * 2 * true).sum(dim=-1) / (true.sum(dim=-1) + 1e-7)
+            ).squeeze()
+        elif self.mode == "regression":
+            rewards = (1 / torch.abs(pred - true)).squeeze()  # b*t
+            rewards = torch.clamp(rewards, min=0, max=5)
+        else:
+            raise ValueError(
+                "mode should be binary, multiclass, multilabel or regression"
+            )
+
+        act_prob1 = model.agent1_prob
+        act_prob1 = torch.stack(act_prob1).permute(1, 0).to(self.device)
+        act_prob1 = act_prob1 * mask.view(act_prob1.size(0), act_prob1.size(1))
+        act_entropy1 = model.agent1_entropy
+        act_entropy1 = torch.stack(act_entropy1).permute(1, 0).to(self.device)
+        act_entropy1 = act_entropy1 * mask.view(
+            act_entropy1.size(0), act_entropy1.size(1)
+        )
+        if self.use_baseline == True:
+            act_baseline1 = model.agent1_baseline
+            act_baseline1 = (
+                torch.stack(act_baseline1).squeeze(-1).permute(1, 0).to(self.device)
+            )
+            act_baseline1 = act_baseline1 * mask.view(
+                act_baseline1.size(0), act_baseline1.size(1)
+            )
+
+        act_prob2 = model.agent2_prob
+        act_prob2 = torch.stack(act_prob2).permute(1, 0).to(self.device)
+        act_prob2 = act_prob2 * mask.view(act_prob2.size(0), act_prob2.size(1))
+        act_entropy2 = model.agent2_entropy
+        act_entropy2 = torch.stack(act_entropy2).permute(1, 0).to(self.device)
+        act_entropy2 = act_entropy2 * mask.view(
+            act_entropy2.size(0), act_entropy2.size(1)
+        )
+        if self.use_baseline == True:
+            act_baseline2 = model.agent2_baseline
+            act_baseline2 = (
+                torch.stack(act_baseline2).squeeze(-1).permute(1, 0).to(self.device)
+            )
+            act_baseline2 = act_baseline2 * mask.view(
+                act_baseline2.size(0), act_baseline2.size(1)
+            )
+
+        running_rewards = []
+        discounted_rewards = 0
+        for i in reversed(range(act_prob1.size(1))):
+            if i == act_prob1.size(1) - 1:
+                discounted_rewards = rewards + gamma * discounted_rewards
+            else:
+                discounted_rewards = (
+                    torch.zeros_like(rewards) + gamma * discounted_rewards
+                )
+            running_rewards.insert(0, discounted_rewards)
+        rewards = torch.stack(running_rewards).permute(1, 0)
+        # rewards = (rewards - rewards.mean(dim=1).unsqueeze(-1)) / (
+        #     rewards.std(dim=1) + 1e-7
+        # ).unsqueeze(-1)
+        rewards = rewards.detach()
+
+        if self.use_baseline == True:
+            loss_value1 = torch.sum((rewards - act_baseline1) ** 2, dim=1) / torch.sum(
+                mask, dim=1
+            )
+            loss_value1 = torch.mean(loss_value1)
+            loss_value2 = torch.sum((rewards - act_baseline2) ** 2, dim=1) / torch.sum(
+                mask, dim=1
+            )
+            loss_value2 = torch.mean(loss_value2)
+            loss_value = loss_value1 + loss_value2
+            loss_RL1 = -torch.sum(
+                act_prob1 * (rewards - act_baseline1) + entropy_term * act_entropy1,
+                dim=1,
+            ) / torch.sum(mask, dim=1)
+            loss_RL1 = torch.mean(loss_RL1)
+            loss_RL2 = -torch.sum(
+                act_prob2 * (rewards - act_baseline2) + entropy_term * act_entropy2,
+                dim=1,
+            ) / torch.sum(mask, dim=1)
+            loss_RL2 = torch.mean(loss_RL2)
+            loss_RL = loss_RL1 + loss_RL2
+            loss = loss_RL + loss_value
+        else:
+            loss_RL1 = -torch.sum(
+                act_prob1 * rewards + entropy_term * act_entropy1, dim=1
+            ) / torch.sum(mask, dim=1)
+            loss_RL1 = torch.mean(loss_RL1)
+            loss_RL2 = -torch.sum(
+                act_prob2 * rewards + entropy_term * act_entropy2, dim=1
+            ) / torch.sum(mask, dim=1)
+            loss_RL2 = torch.mean(loss_RL2)
+            loss_RL = loss_RL1 + loss_RL2
+            loss = loss_RL
+
+        return loss
+
+    def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
+        """Forward propagation.
+
+        The label `kwargs[self.label_key]` is a list of labels for each patient.
+
+        Args:
+            **kwargs: keyword arguments for the model. The keys must contain
+                all the feature keys and the label key.
+
+        Returns:
+            A dictionary with the following keys:
+                loss: a scalar tensor representing the final loss.
+                loss_task: a scalar tensor representing the task loss.
+                loss_RL: a scalar tensor representing the RL loss.
+                y_prob: a tensor representing the predicted probabilities.
+                y_true: a tensor representing the true labels.
+        """
+        patient_emb = []
+        mask_dict = {}
+        for feature_key in self.feature_keys:
+            input_info = self.dataset.input_info[feature_key]
+            dim_, type_ = input_info["dim"], input_info["type"]
+
+            # for case 1: [code1, code2, code3, ...]
+            if (dim_ == 2) and (type_ == str):
+                x = self.feat_tokenizers[feature_key].batch_encode_2d(
+                    kwargs[feature_key]
+                )
+                # (patient, event)
+                x = torch.tensor(x, dtype=torch.long, device=self.device)
+                # (patient, event, embedding_dim)
+                x = self.embeddings[feature_key](x)
+                # (patient, event)
+                mask = torch.sum(x, dim=2) != 0
+                mask_dict[feature_key] = mask
+
+            # for case 2: [[code1, code2], [code3, ...], ...]
+            elif (dim_ == 3) and (type_ == str):
+                x = self.feat_tokenizers[feature_key].batch_encode_3d(
+                    kwargs[feature_key]
+                )
+                # (patient, visit, event)
+                x = torch.tensor(x, dtype=torch.long, device=self.device)
+                # (patient, visit, event, embedding_dim)
+                x = self.embeddings[feature_key](x)
+                # (patient, visit, embedding_dim)
+                x = torch.sum(x, dim=2)
+                # (patient, visit)
+                mask = torch.sum(x, dim=2) != 0
+                mask_dict[feature_key] = mask
+
+            # for case 3: [[1.5, 2.0, 0.0], ...]
+            elif (dim_ == 2) and (type_ in [float, int]):
+                x, mask = self.padding2d(kwargs[feature_key])
+                # (patient, event, values)
+                x = torch.tensor(x, dtype=torch.float, device=self.device)
+                # (patient, event, embedding_dim)
+                x = self.linear_layers[feature_key](x)
+                # (patient, event)
+                mask = mask.bool().to(self.device)
+                mask_dict[feature_key] = mask
+
+            # for case 4: [[[1.5, 2.0, 0.0], [1.8, 2.4, 6.0]], ...]
+            elif (dim_ == 3) and (type_ in [float, int]):
+                x, mask = self.padding3d(kwargs[feature_key])
+                # (patient, visit, event, values)
+                x = torch.tensor(x, dtype=torch.float, device=self.device)
+                # (patient, visit, embedding_dim)
+                x = torch.sum(x, dim=2)
+                x = self.linear_layers[feature_key](x)
+                # (patient, event)
+                mask = mask[:, :, 0]
+                mask = mask.bool().to(self.device)
+                mask_dict[feature_key] = mask
+
+            else:
+                raise NotImplementedError
+
+            if self.static_dim > 0:
+                static = torch.tensor(
+                    kwargs[self.static_key], dtype=torch.float, device=self.device
+                )
+                x, _ = self.agent[feature_key](x, static=static, mask=mask)
+            else:
+                x, _ = self.agent[feature_key](x, mask=mask)
+            patient_emb.append(x)
+
+        patient_emb = torch.cat(patient_emb, dim=1)
+        # (patient, label_size)
+        logits = self.fc(patient_emb)
+        # obtain y_true, loss, y_prob
+        y_true = self.prepare_labels(kwargs[self.label_key], self.label_tokenizer)
+        loss_task = self.get_loss_function()(logits, y_true)
+
+        loss_rl = 0
+        for feature_key in self.feature_keys:
+            cur_loss = self.get_loss(
+                self.agent[feature_key], logits, y_true, mask_dict[feature_key]
+            )
+            loss_rl += cur_loss
+        loss = loss_task + loss_rl
+        y_prob = self.prepare_y_prob(logits)
+        # return {
+        #     "loss": loss,
+        #     "loss_task": loss_task,
+        #     "loss_rl": loss_rl,
+        #     "y_prob": y_prob,
+        #     "y_true": y_true,
+        # }
+        return {
+            "loss": loss,
+            "y_prob": y_prob,
+            "y_true": y_true,
+        }
+
+
+if __name__ == "__main__":
+    from pyhealth.datasets import SampleEHRDataset
+
+    samples = [
+        {
+            "patient_id": "patient-0",
+            "visit_id": "visit-0",
+            # "single_vector": [1, 2, 3],
+            "list_codes": ["505800458", "50580045810", "50580045811"],  # NDC
+            "list_vectors": [[1.0, 2.55, 3.4], [4.1, 5.5, 6.0]],
+            "list_list_codes": [["A05B", "A05C", "A06A"], ["A11D", "A11E"]],  # ATC-4
+            "list_list_vectors": [
+                [[1.8, 2.25, 3.41], [4.50, 5.9, 6.0]],
+                [[7.7, 8.5, 9.4]],
+            ],
+            "label": 1,
+            "demographic": [1.0, 2.0, 1.3],
+        },
+        {
+            "patient_id": "patient-0",
+            "visit_id": "visit-1",
+            # "single_vector": [1, 5, 8],
+            "list_codes": [
+                "55154191800",
+                "551541928",
+                "55154192800",
+                "705182798",
+                "70518279800",
+            ],
+            "list_vectors": [[1.4, 3.2, 3.5], [4.1, 5.9, 1.7], [4.5, 5.9, 1.7]],
+            "list_list_codes": [["A04A", "B035", "C129"]],
+            "list_list_vectors": [
+                [[1.0, 2.8, 3.3], [4.9, 5.0, 6.6], [7.7, 8.4, 1.3], [7.7, 8.4, 1.3]],
+            ],
+            "label": 0,
+            "demographic": [1.0, 2.0, 1.3],
+        },
+    ]
+
+    # dataset
+    dataset = SampleEHRDataset(samples=samples, dataset_name="test")
+
+    # data loader
+    from pyhealth.datasets import get_dataloader
+
+    train_loader = get_dataloader(dataset, batch_size=2, shuffle=True)
+
+    # model
+    model = Agent(
+        dataset=dataset,
+        feature_keys=[
+            "list_codes",
+            "list_vectors",
+            "list_list_codes",
+            # "list_list_vectors",
+        ],
+        static_key="demographic",
+        label_key="label",
+        mode="binary",
+    )
+
+    # data batch
+    data_batch = next(iter(train_loader))
+
+    # try the model
+    ret = model(**data_batch)
+    print(ret)
+
+    # try loss backward
+    ret["loss"].backward()
diff --git a/pyhealth/models/concare.py b/pyhealth/models/concare.py
new file mode 100644
index 00000000..1137acd0
--- /dev/null
+++ b/pyhealth/models/concare.py
@@ -0,0 +1,1007 @@
+from typing import List, Tuple, Dict, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.utils.rnn as rnn_utils
+
+from pyhealth.datasets import SampleEHRDataset
+from pyhealth.models import BaseModel
+from pyhealth.models.utils import get_last_visit
+
+import math
+
+
+class FinalAttentionQKV(nn.Module):
+    def __init__(
+        self,
+        attention_input_dim: int,
+        attention_hidden_dim: int,
+        attention_type: str = "add",
+        dropout: float = 0.5,
+    ):
+        super(FinalAttentionQKV, self).__init__()
+
+        self.attention_type = attention_type
+        self.attention_hidden_dim = attention_hidden_dim
+        self.attention_input_dim = attention_input_dim
+
+        self.W_q = nn.Linear(attention_input_dim, attention_hidden_dim)
+        self.W_k = nn.Linear(attention_input_dim, attention_hidden_dim)
+        self.W_v = nn.Linear(attention_input_dim, attention_hidden_dim)
+
+        self.W_out = nn.Linear(attention_hidden_dim, 1)
+
+        self.b_in = nn.Parameter(
+            torch.zeros(
+                1,
+            )
+        )
+        self.b_out = nn.Parameter(
+            torch.zeros(
+                1,
+            )
+        )
+
+        nn.init.kaiming_uniform_(self.W_q.weight, a=math.sqrt(5))
+        nn.init.kaiming_uniform_(self.W_k.weight, a=math.sqrt(5))
+        nn.init.kaiming_uniform_(self.W_v.weight, a=math.sqrt(5))
+        nn.init.kaiming_uniform_(self.W_out.weight, a=math.sqrt(5))
+
+        self.Wh = nn.Parameter(
+            torch.randn(2 * attention_input_dim, attention_hidden_dim)
+        )
+        self.Wa = nn.Parameter(torch.randn(attention_hidden_dim, 1))
+        self.ba = nn.Parameter(
+            torch.zeros(
+                1,
+            )
+        )
+
+        nn.init.kaiming_uniform_(self.Wh, a=math.sqrt(5))
+        nn.init.kaiming_uniform_(self.Wa, a=math.sqrt(5))
+
+        self.dropout = nn.Dropout(p=dropout)
+        self.tanh = nn.Tanh()
+        self.softmax = nn.Softmax(dim=1)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, input):
+
+        (
+            batch_size,
+            time_step,
+            input_dim,
+        ) = input.size()  # batch_size * input_dim + 1 * hidden_dim(i)
+        input_q = self.W_q(input[:, -1, :])  # b h
+        input_k = self.W_k(input)  # b t h
+        input_v = self.W_v(input)  # b t h
+
+        if self.attention_type == "add":  # B*T*I  @ H*I
+
+            q = torch.reshape(
+                input_q, (batch_size, 1, self.attention_hidden_dim)
+            )  # B*1*H
+            h = q + input_k + self.b_in  # b t h
+            h = self.tanh(h)  # B*T*H
+            e = self.W_out(h)  # b t 1
+            e = torch.reshape(e, (batch_size, time_step))  # b t
+
+        elif self.attention_type == "mul":
+            q = torch.reshape(
+                input_q, (batch_size, self.attention_hidden_dim, 1)
+            )  # B*h 1
+            e = torch.matmul(input_k, q).squeeze()  # b t
+
+        elif self.attention_type == "concat":
+            q = input_q.unsqueeze(1).repeat(1, time_step, 1)  # b t h
+            k = input_k
+            c = torch.cat((q, k), dim=-1)  # B*T*2I
+            h = torch.matmul(c, self.Wh)
+            h = self.tanh(h)
+            e = torch.matmul(h, self.Wa) + self.ba  # B*T*1
+            e = torch.reshape(e, (batch_size, time_step))  # b t
+        else:
+            raise ValueError(
+                "Unknown attention type: {}, please use add, mul, concat".format(
+                    self.attention_type
+                )
+            )
+
+        a = self.softmax(e)  # B*T
+        if self.dropout is not None:
+            a = self.dropout(a)
+        v = torch.matmul(a.unsqueeze(1), input_v).squeeze()  # B*I
+
+        return v, a
+
+
+class PositionwiseFeedForward(nn.Module):  # new added
+    "Implements FFN equation."
+
+    def __init__(self, d_model, d_ff, dropout=0.1):
+        super(PositionwiseFeedForward, self).__init__()
+        self.w_1 = nn.Linear(d_model, d_ff)
+        self.w_2 = nn.Linear(d_ff, d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        return self.w_2(self.dropout(torch.relu(self.w_1(x)))), None
+
+
+class PositionalEncoding(nn.Module):  # new added / not use anymore
+    "Implement the PE function."
+
+    def __init__(self, d_model, dropout, max_len=400):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        # Compute the positional encodings once in log space.
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0.0, max_len).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0.0, d_model, 2) * -(math.log(10000.0) / d_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        pos = self.pe[:, : x.size(1)].clone().requires_grad_(False)
+        x = x + pos
+        return self.dropout(x)
+
+
+class MultiHeadedAttention(nn.Module):
+    def __init__(self, h, d_model, dropout=0):
+        "Take in model size and number of heads."
+        super(MultiHeadedAttention, self).__init__()
+        assert d_model % h == 0
+        # We assume d_v always equals d_k
+        self.d_k = d_model // h
+        self.h = h
+        self.linears = nn.ModuleList(
+            [nn.Linear(d_model, self.d_k * self.h) for _ in range(3)]
+        )
+        self.final_linear = nn.Linear(d_model, d_model)
+        self.attn = None
+        self.dropout = nn.Dropout(p=dropout)
+
+    def attention(self, query, key, value, mask=None, dropout=None):
+        "Compute 'Scaled Dot Product Attention'"
+        d_k = query.size(-1)  # b h t d_k
+        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)  # b h t t
+        if mask is not None:  # 1 1 t t
+            scores = scores.masked_fill(mask == 0, -1e9)  # b h t t 下三角
+        p_attn = torch.softmax(scores, dim=-1)  # b h t t
+        if dropout is not None:
+            p_attn = dropout(p_attn)
+        return torch.matmul(p_attn, value), p_attn  # b h t v (d_k)
+
+    def cov(self, m, y=None):
+        if y is not None:
+            m = torch.cat((m, y), dim=0)
+        m_exp = torch.mean(m, dim=1)
+        x = m - m_exp[:, None]
+        cov = 1 / (x.size(1) - 1) * x.mm(x.t())
+        return cov
+
+    def forward(self, query, key, value, mask=None):
+        if mask is not None:
+            # Same mask applied to all h heads.
+            mask = mask.unsqueeze(1)  # 1 1 t t
+
+        nbatches = query.size(0)  # b
+        input_dim = query.size(1)  # i+1
+        feature_dim = query.size(1)  # i+1
+
+        # input size -> # batch_size * d_input * hidden_dim
+
+        # d_model => h * d_k
+        query, key, value = [
+            l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
+            for l, x in zip(self.linears, (query, key, value))
+        ]  # b num_head d_input d_k
+
+        x, self.attn = self.attention(
+            query, key, value, mask=mask, dropout=self.dropout
+        )  # b num_head d_input d_v (d_k)
+
+        x = (
+            x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)
+        )  # batch_size * d_input * hidden_dim
+
+        # DeCov
+        DeCov_contexts = x.transpose(0, 1).transpose(1, 2)  # I+1 H B
+        Covs = self.cov(DeCov_contexts[0, :, :])
+        DeCov_loss = 0.5 * (
+            torch.norm(Covs, p="fro") ** 2 - torch.norm(torch.diag(Covs)) ** 2
+        )
+        for i in range(feature_dim - 1):
+            Covs = self.cov(DeCov_contexts[i + 1, :, :])
+            DeCov_loss += 0.5 * (
+                torch.norm(Covs, p="fro") ** 2 - torch.norm(torch.diag(Covs)) ** 2
+            )
+
+        return self.final_linear(x), DeCov_loss
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, features, eps=1e-7):
+        super(LayerNorm, self).__init__()
+        self.a_2 = nn.Parameter(torch.ones(features))
+        self.b_2 = nn.Parameter(torch.zeros(features))
+        self.eps = eps
+
+    def forward(self, x):
+        mean = x.mean(-1, keepdim=True)
+        std = x.std(-1, keepdim=True)
+        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
+
+
+class SublayerConnection(nn.Module):
+    """
+    A residual connection followed by a layer norm.
+    Note for code simplicity the norm is first as opposed to last.
+    """
+
+    def __init__(self, size, dropout):
+        super(SublayerConnection, self).__init__()
+        self.norm = LayerNorm(size)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, sublayer):
+        "Apply residual connection to any sublayer with the same size."
+        returned_value = sublayer(self.norm(x))
+        return x + self.dropout(returned_value[0]), returned_value[1]
+
+
+class SingleAttention(nn.Module):
+    def __init__(
+        self,
+        attention_input_dim,
+        attention_hidden_dim,
+        attention_type="add",
+        time_aware=False,
+    ):
+        super(SingleAttention, self).__init__()
+
+        self.attention_type = attention_type
+        self.attention_hidden_dim = attention_hidden_dim
+        self.attention_input_dim = attention_input_dim
+        self.time_aware = time_aware
+
+        # batch_time = torch.arange(0, batch_mask.size()[1], dtype=torch.float32).reshape(1, batch_mask.size()[1], 1)
+        # batch_time = batch_time.repeat(batch_mask.size()[0], 1, 1)
+
+        if attention_type == "add":
+            if self.time_aware:
+                # self.Wx = nn.Parameter(torch.randn(attention_input_dim+1, attention_hidden_dim))
+                self.Wx = nn.Parameter(
+                    torch.randn(attention_input_dim, attention_hidden_dim)
+                )
+                self.Wtime_aware = nn.Parameter(torch.randn(1, attention_hidden_dim))
+                nn.init.kaiming_uniform_(self.Wtime_aware, a=math.sqrt(5))
+            else:
+                self.Wx = nn.Parameter(
+                    torch.randn(attention_input_dim, attention_hidden_dim)
+                )
+            self.Wt = nn.Parameter(
+                torch.randn(attention_input_dim, attention_hidden_dim)
+            )
+            self.bh = nn.Parameter(
+                torch.zeros(
+                    attention_hidden_dim,
+                )
+            )
+            self.Wa = nn.Parameter(torch.randn(attention_hidden_dim, 1))
+            self.ba = nn.Parameter(
+                torch.zeros(
+                    1,
+                )
+            )
+
+            nn.init.kaiming_uniform_(self.Wd, a=math.sqrt(5))
+            nn.init.kaiming_uniform_(self.Wx, a=math.sqrt(5))
+            nn.init.kaiming_uniform_(self.Wt, a=math.sqrt(5))
+            nn.init.kaiming_uniform_(self.Wa, a=math.sqrt(5))
+        elif attention_type == "mul":
+            self.Wa = nn.Parameter(
+                torch.randn(attention_input_dim, attention_input_dim)
+            )
+            self.ba = nn.Parameter(
+                torch.zeros(
+                    1,
+                )
+            )
+
+            nn.init.kaiming_uniform_(self.Wa, a=math.sqrt(5))
+        elif attention_type == "concat":
+            if self.time_aware:
+                self.Wh = nn.Parameter(
+                    torch.randn(2 * attention_input_dim + 1, attention_hidden_dim)
+                )
+            else:
+                self.Wh = nn.Parameter(
+                    torch.randn(2 * attention_input_dim, attention_hidden_dim)
+                )
+
+            self.Wa = nn.Parameter(torch.randn(attention_hidden_dim, 1))
+            self.ba = nn.Parameter(
+                torch.zeros(
+                    1,
+                )
+            )
+
+            nn.init.kaiming_uniform_(self.Wh, a=math.sqrt(5))
+            nn.init.kaiming_uniform_(self.Wa, a=math.sqrt(5))
+
+        elif attention_type == "new":
+            self.Wt = nn.Parameter(
+                torch.randn(attention_input_dim, attention_hidden_dim)
+            )
+            self.Wx = nn.Parameter(
+                torch.randn(attention_input_dim, attention_hidden_dim)
+            )
+
+            self.rate = nn.Parameter(torch.zeros(1) + 0.8)
+            nn.init.kaiming_uniform_(self.Wx, a=math.sqrt(5))
+            nn.init.kaiming_uniform_(self.Wt, a=math.sqrt(5))
+
+        else:
+            raise RuntimeError(
+                "Wrong attention type. Please use 'add', 'mul', 'concat' or 'new'."
+            )
+
+        self.tanh = nn.Tanh()
+        self.softmax = nn.Softmax(dim=1)
+        self.sigmoid = nn.Sigmoid()
+        self.relu = nn.ReLU()
+
+    def forward(self, input, mask, device):
+
+        (
+            batch_size,
+            time_step,
+            input_dim,
+        ) = input.size()  # batch_size * time_step * hidden_dim(i)
+
+        time_decays = (
+            torch.tensor(range(time_step - 1, -1, -1), dtype=torch.float32)
+            .unsqueeze(-1)
+            .unsqueeze(0)
+            .to(device=device)
+        )  # 1*t*1
+        b_time_decays = time_decays.repeat(batch_size, 1, 1) + 1  # b t 1
+
+        if self.attention_type == "add":  # B*T*I  @ H*I
+            last_visit = get_last_visit(input, mask)
+            q = torch.matmul(last_visit, self.Wt)  # b h
+            q = torch.reshape(q, (batch_size, 1, self.attention_hidden_dim))  # B*1*H
+            if self.time_aware == True:
+                k = torch.matmul(input, self.Wx)  # b t h
+                time_hidden = torch.matmul(b_time_decays, self.Wtime_aware)  # b t h
+            else:
+                k = torch.matmul(input, self.Wx)  # b t h
+            h = q + k + self.bh  # b t h
+            if self.time_aware:
+                h += time_hidden
+            h = self.tanh(h)  # B*T*H
+            e = torch.matmul(h, self.Wa) + self.ba  # B*T*1
+            e = torch.reshape(e, (batch_size, time_step))  # b t
+        elif self.attention_type == "mul":
+            last_visit = get_last_visit(input, mask)
+            e = torch.matmul(last_visit, self.Wa)  # b i
+            e = (
+                torch.matmul(e.unsqueeze(1), input.permute(0, 2, 1)).reshape(
+                    batch_size, time_step
+                )
+                + self.ba
+            )  # b t
+        elif self.attention_type == "concat":
+            last_visit = get_last_visit(input, mask)
+            q = last_visit.unsqueeze(1).repeat(1, time_step, 1)  # b t i
+            k = input
+            c = torch.cat((q, k), dim=-1)  # B*T*2I
+            if self.time_aware:
+                c = torch.cat((c, b_time_decays), dim=-1)  # B*T*2I+1
+            h = torch.matmul(c, self.Wh)
+            h = self.tanh(h)
+            e = torch.matmul(h, self.Wa) + self.ba  # B*T*1
+            e = torch.reshape(e, (batch_size, time_step))  # b t
+
+        elif self.attention_type == "new":
+            last_visit = get_last_visit(input, mask)
+            q = torch.matmul(last_visit, self.Wt)  # b h
+            q = torch.reshape(q, (batch_size, 1, self.attention_hidden_dim))  # B*1*H
+            k = torch.matmul(input, self.Wx)  # b t h
+            dot_product = torch.matmul(q, k.transpose(1, 2)).reshape(
+                batch_size, time_step
+            )  # b t
+            denominator = self.sigmoid(self.rate) * (
+                torch.log(2.72 + (1 - self.sigmoid(dot_product)))
+                * (b_time_decays.reshape(batch_size, time_step))
+            )
+            e = self.relu(self.sigmoid(dot_product) / (denominator))  # b * t
+        else:
+            raise ValueError(
+                "Wrong attention type. Plase use 'add', 'mul', 'concat' or 'new'."
+            )
+
+        if mask is not None:
+            e = e.masked_fill(mask == 0, -1e9)
+        a = self.softmax(e)  # B*T
+        v = torch.matmul(a.unsqueeze(1), input).reshape(batch_size, input_dim)  # B*I
+
+        return v, a
+
+
+class ConCareLayer(nn.Module):
+    """ConCare layer.
+
+    Paper: Liantao Ma et al. Concare: Personalized clinical feature embedding via capturing the healthcare context. AAAI 2020.
+
+    This layer is used in the ConCare model. But it can also be used as a
+    standalone layer.
+
+    Args:
+        input_dim: dynamic feature size.
+        static_dim: static feature size, if 0, then no static feature is used.
+        hidden_dim: hidden dimension of the channel-wise GRU, default 128.
+        transformer_hidden: hidden dimension of the transformer, default 128.
+        num_head: number of heads in the transformer, default 4.
+        pe_hidden: hidden dimension of the positional encoding, default 64.
+        dropout: dropout rate, default 0.5.
+
+    Examples:
+        >>> from pyhealth.models import ConCareLayer
+        >>> input = torch.randn(3, 128, 64)  # [batch size, sequence len, feature_size]
+        >>> layer = ConCareLayer(64)
+        >>> c, _ = layer(input)
+        >>> c.shape
+        torch.Size([3, 128])
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        static_dim: int = 0,
+        hidden_dim: int = 128,
+        num_head: int = 4,
+        pe_hidden: int = 64,
+        dropout: int = 0.5,
+    ):
+        super(ConCareLayer, self).__init__()
+
+        # hyperparameters
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim  # d_model
+        self.transformer_hidden = hidden_dim
+        self.num_head = num_head
+        self.pe_hidden = pe_hidden
+        # self.output_dim = output_dim
+        self.dropout = dropout
+        self.static_dim = static_dim
+
+        # layers
+        self.PositionalEncoding = PositionalEncoding(
+            self.transformer_hidden, dropout=0, max_len=400
+        )
+
+        self.GRUs = nn.ModuleList(
+            [
+                nn.GRU(1, self.hidden_dim, batch_first=True)
+                for _ in range(self.input_dim)
+            ]
+        )
+        self.LastStepAttentions = nn.ModuleList(
+            [
+                SingleAttention(
+                    self.hidden_dim,
+                    8,
+                    attention_type="new",
+                    time_aware=True,
+                )
+                for _ in range(self.input_dim)
+            ]
+        )
+
+        self.FinalAttentionQKV = FinalAttentionQKV(
+            self.hidden_dim,
+            self.hidden_dim,
+            attention_type="mul",
+            dropout=self.dropout,
+        )
+
+        self.MultiHeadedAttention = MultiHeadedAttention(
+            self.num_head, self.transformer_hidden, dropout=self.dropout
+        )
+        self.SublayerConnection = SublayerConnection(
+            self.transformer_hidden, dropout=self.dropout
+        )
+
+        self.PositionwiseFeedForward = PositionwiseFeedForward(
+            self.transformer_hidden, self.pe_hidden, dropout=0.1
+        )
+
+        if self.static_dim > 0:
+            self.demo_proj_main = nn.Linear(self.static_dim, self.hidden_dim)
+
+        self.dropout = nn.Dropout(p=self.dropout)
+        self.tanh = nn.Tanh()
+        self.softmax = nn.Softmax()
+        self.sigmoid = nn.Sigmoid()
+        self.relu = nn.ReLU()
+
+    def concare_encoder(self, input, static=None, mask=None):
+        # input shape [batch_size, timestep, feature_dim]
+
+        if self.static_dim > 0:
+            demo_main = self.tanh(self.demo_proj_main(static)).unsqueeze(
+                1
+            )  # b hidden_dim
+
+        batch_size = input.size(0)
+        time_step = input.size(1)
+        feature_dim = input.size(2)
+
+        if self.transformer_hidden % self.num_head != 0:
+            raise ValueError("transformer_hidden must be divisible by num_head")
+
+        # forward
+        GRU_embeded_input = self.GRUs[0](
+            input[:, :, 0].unsqueeze(-1).to(device=input.device),
+            torch.zeros(batch_size, self.hidden_dim)
+            .to(device=input.device)
+            .unsqueeze(0),
+        )[
+            0
+        ]  # b t h
+        Attention_embeded_input = self.LastStepAttentions[0](
+            GRU_embeded_input, mask, input.device
+        )[0].unsqueeze(
+            1
+        )  # b 1 h
+
+        for i in range(feature_dim - 1):
+            embeded_input = self.GRUs[i + 1](
+                input[:, :, i + 1].unsqueeze(-1),
+                torch.zeros(batch_size, self.hidden_dim)
+                .to(device=input.device)
+                .unsqueeze(0),
+            )[
+                0
+            ]  # b 1 h
+            embeded_input = self.LastStepAttentions[i + 1](
+                embeded_input, mask, input.device
+            )[0].unsqueeze(
+                1
+            )  # b 1 h
+            Attention_embeded_input = torch.cat(
+                (Attention_embeded_input, embeded_input), 1
+            )  # b i h
+
+        if self.static_dim > 0:
+            Attention_embeded_input = torch.cat(
+                (Attention_embeded_input, demo_main), 1
+            )  # b i+1 h
+        posi_input = self.dropout(
+            Attention_embeded_input
+        )  # batch_size * d_input+1 * hidden_dim
+
+        contexts = self.SublayerConnection(
+            posi_input,
+            lambda x: self.MultiHeadedAttention(
+                posi_input, posi_input, posi_input, None
+            ),
+        )  # # batch_size * d_input * hidden_dim
+
+        DeCov_loss = contexts[1]
+        contexts = contexts[0]
+
+        contexts = self.SublayerConnection(
+            contexts, lambda x: self.PositionwiseFeedForward(contexts)
+        )[0]
+
+        weighted_contexts, a = self.FinalAttentionQKV(contexts)
+        return weighted_contexts, DeCov_loss
+
+    def forward(
+        self,
+        x: torch.tensor,
+        static: Optional[torch.tensor] = None,
+        mask: Optional[torch.tensor] = None,
+    ) -> Tuple[torch.tensor]:
+        """Forward propagation.
+
+        Args:
+            x: a tensor of shape [batch size, sequence len, input_dim].
+            static: a tensor of shape [batch size, static_dim].
+            mask: an optional tensor of shape [batch size, sequence len], where
+                1 indicates valid and 0 indicates invalid.
+
+        Returns:
+            output: a tensor of shape [batch size, fusion_dim] representing the
+                patient embedding.
+            decov: the decov loss value
+        """
+        # rnn will only apply dropout between layers
+        batch_size, time_steps, _ = x.size()
+        out = torch.zeros((batch_size, time_steps, self.hidden_dim))
+        out, decov = self.concare_encoder(x, static, mask)
+        out = self.dropout(out)
+        return out, decov
+
+
+class ConCare(BaseModel):
+    """ConCare model.
+
+    Paper: Liantao Ma et al. Concare: Personalized clinical feature embedding via capturing the healthcare context. AAAI 2020.
+
+    Note:
+        We use separate ConCare layers for different feature_keys.
+        Currently, we automatically support different input formats:
+            - code based input (need to use the embedding table later)
+            - float/int based value input
+        If you need the interpretable feature correlations provided by the ConCare model calculates the , we do not recommend use embeddings for the input features.
+        We follow the current convention for the ConCare model:
+            - case 1. [code1, code2, code3, ...]
+                - we will assume the code follows the order; our model will encode
+                each code into a vector and apply ConCare on the code level
+            - case 2. [[code1, code2]] or [[code1, code2], [code3, code4, code5], ...]
+                - we will assume the inner bracket follows the order; our model first
+                use the embedding table to encode each code into a vector and then use
+                average/mean pooling to get one vector for one inner bracket; then use
+                ConCare one the braket level
+            - case 3. [[1.5, 2.0, 0.0]] or [[1.5, 2.0, 0.0], [8, 1.2, 4.5], ...]
+                - this case only makes sense when each inner bracket has the same length;
+                we assume each dimension has the same meaning; we run ConCare directly
+                on the inner bracket level, similar to case 1 after embedding table
+            - case 4. [[[1.5, 2.0, 0.0]]] or [[[1.5, 2.0, 0.0], [8, 1.2, 4.5]], ...]
+                - this case only makes sense when each inner bracket has the same length;
+                we assume each dimension has the same meaning; we run ConCare directly
+                on the inner bracket level, similar to case 2 after embedding table
+
+    Args:
+        dataset: the dataset to train the model. It is used to query certain
+            information such as the set of all tokens.
+        feature_keys:  list of keys in samples to use as features,
+            e.g. ["conditions", "procedures"].
+        label_key: key in samples to use as label (e.g., "drugs").
+        mode: one of "binary", "multiclass", or "multilabel".
+        static_keys: the key in samples to use as static features, e.g. "demographics". Default is None.
+                     we only support numerical static features.
+        use_embedding: list of bools indicating whether to use embedding for each feature type,
+            e.g. [True, False].
+        embedding_dim: the embedding dimension. Default is 128.
+        hidden_dim: the hidden dimension. Default is 128.
+        **kwargs: other parameters for the ConCare layer.
+
+
+    Examples:
+        >>> from pyhealth.datasets import SampleEHRDataset
+        >>> samples = [
+        ...         {
+        ...             "patient_id": "patient-0",
+        ...             "visit_id": "visit-0",
+        ...             "list_codes": ["505800458", "50580045810", "50580045811"],  # NDC
+        ...             "list_vectors": [[1.0, 2.55, 3.4], [4.1, 5.5, 6.0]],
+        ...             "list_list_codes": [["A05B", "A05C", "A06A"], ["A11D", "A11E"]],  # ATC-4
+        ...             "list_list_vectors": [
+        ...                 [[1.8, 2.25, 3.41], [4.50, 5.9, 6.0]],
+        ...                 [[7.7, 8.5, 9.4]],
+        ...             ],
+        ...             "demographic": [0.0, 2.0, 1.5],
+        ...             "label": 1,
+        ...         },
+        ...         {
+        ...             "patient_id": "patient-0",
+        ...             "visit_id": "visit-1",
+        ...             "list_codes": [
+        ...                 "55154191800",
+        ...                 "551541928",
+        ...                 "55154192800",
+        ...                 "705182798",
+        ...                 "70518279800",
+        ...             ],
+        ...             "list_vectors": [[1.4, 3.2, 3.5], [4.1, 5.9, 1.7], [4.5, 5.9, 1.7]],
+        ...             "list_list_codes": [["A04A", "B035", "C129"]],
+        ...             "list_list_vectors": [
+        ...                 [[1.0, 2.8, 3.3], [4.9, 5.0, 6.6], [7.7, 8.4, 1.3], [7.7, 8.4, 1.3]],
+        ...             ],
+        ...             "demographic": [0.0, 2.0, 1.5],
+        ...             "label": 0,
+        ...         },
+        ...     ]
+        >>> dataset = SampleEHRDataset(samples=samples, dataset_name="test")
+        >>>
+        >>> from pyhealth.models import ConCare
+        >>> model = ConCare(
+        ...         dataset=dataset,
+        ...         feature_keys=[
+        ...             "list_codes",
+        ...             "list_vectors",
+        ...             "list_list_codes",
+        ...             "list_list_vectors",
+        ...         ],
+        ...         label_key="label",
+        ...         static_key="demographic",
+        ...         use_embedding=[True, False, True, False],
+        ...         mode="binary"
+        ...     )
+        >>>
+        >>> from pyhealth.datasets import get_dataloader
+        >>> train_loader = get_dataloader(dataset, batch_size=2, shuffle=True)
+        >>> data_batch = next(iter(train_loader))
+        >>>
+        >>> ret = model(**data_batch)
+        >>> print(ret)
+        {'loss': tensor(28.2995, grad_fn=<AddBackward0>), 'loss_task': tensor(0.7120, grad_fn=<BinaryCrossEntropyWithLogitsBackward0>), 'decov_loss': tensor(27.5875, grad_fn=<AddBackward0>), 'y_prob': tensor([[0.5389], [0.5533]], grad_fn=<SigmoidBackward0>), 'y_true': tensor([[1.], [0.]])}
+        >>>
+
+    """
+
+    def __init__(
+        self,
+        dataset: SampleEHRDataset,
+        feature_keys: List[str],
+        label_key: str,
+        mode: str,
+        use_embedding: List[bool],
+        static_key: Optional[str] = None,
+        embedding_dim: int = 128,
+        hidden_dim: int = 128,
+        **kwargs,
+    ):
+        super(ConCare, self).__init__(
+            dataset=dataset,
+            feature_keys=feature_keys,
+            label_key=label_key,
+            mode=mode,
+        )
+        self.embedding_dim = embedding_dim
+        self.use_embedding = use_embedding
+        self.hidden_dim = hidden_dim
+
+        # validate kwargs for ConCare layer
+        if "feature_size" in kwargs:
+            raise ValueError("feature_size is determined by embedding_dim")
+
+        # the key of self.feat_tokenizers only contains the code based inputs
+        self.feat_tokenizers = {}
+        self.static_key = static_key
+        self.label_tokenizer = self.get_label_tokenizer()
+        # the key of self.embeddings only contains the code based inputs
+        self.embeddings = nn.ModuleDict()
+        # the key of self.linear_layers only contains the float/int based inputs
+        self.linear_layers = nn.ModuleDict()
+
+        self.static_dim = 0
+        if self.static_key is not None:
+            self.static_dim = self.dataset.input_info[self.static_key]["len"]
+
+        self.concare = nn.ModuleDict()
+        # add feature ConCare layers
+        for idx, feature_key in enumerate(self.feature_keys):
+            input_info = self.dataset.input_info[feature_key]
+            # sanity check
+            if input_info["type"] not in [str, float, int]:
+                raise ValueError(
+                    "ConCare only supports str code, float and int as input types"
+                )
+            elif (input_info["type"] == str) and (input_info["dim"] not in [2, 3]):
+                raise ValueError(
+                    "ConCare only supports 2-dim or 3-dim str code as input types"
+                )
+            elif (input_info["type"] == str) and (use_embedding[idx] == False):
+                raise ValueError(
+                    "ConCare only supports embedding for str code as input types"
+                )
+            elif (input_info["type"] in [float, int]) and (
+                input_info["dim"] not in [2, 3]
+            ):
+                raise ValueError(
+                    "ConCare only supports 2-dim or 3-dim float and int as input types"
+                )
+
+            # for code based input, we need Type
+            # for float/int based input, we need Type, input_dim
+            if use_embedding[idx]:
+                self.add_feature_transform_layer(feature_key, input_info)
+                self.concare[feature_key] = ConCareLayer(
+                    input_dim=embedding_dim,
+                    static_dim=self.static_dim,
+                    hidden_dim=self.hidden_dim,
+                    **kwargs,
+                )
+            else:
+                self.concare[feature_key] = ConCareLayer(
+                    input_dim=input_info["len"],
+                    static_dim=self.static_dim,
+                    hidden_dim=self.hidden_dim,
+                    **kwargs,
+                )
+
+        output_size = self.get_output_size(self.label_tokenizer)
+        self.fc = nn.Linear(len(self.feature_keys) * self.hidden_dim, output_size)
+
+    def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
+        """Forward propagation.
+
+        The label `kwargs[self.label_key]` is a list of labels for each patient.
+
+        Args:
+            **kwargs: keyword arguments for the model. The keys must contain
+                all the feature keys and the label key.
+
+        Returns:
+            A dictionary with the following keys:
+                loss: a scalar tensor representing the final loss.
+                loss_task: a scalar tensor representing the task loss.
+                loss_decov: a scalar tensor representing the decov loss.
+                y_prob: a tensor representing the predicted probabilities.
+                y_true: a tensor representing the true labels.
+        """
+        patient_emb = []
+        decov_loss = 0
+        for idx, feature_key in enumerate(self.feature_keys):
+            input_info = self.dataset.input_info[feature_key]
+            dim_, type_ = input_info["dim"], input_info["type"]
+
+            # for case 1: [code1, code2, code3, ...]
+            if (dim_ == 2) and (type_ == str):
+                x = self.feat_tokenizers[feature_key].batch_encode_2d(
+                    kwargs[feature_key]
+                )
+                # (patient, event)
+                x = torch.tensor(x, dtype=torch.long, device=self.device)
+                # (patient, event, embedding_dim)
+                x = self.embeddings[feature_key](x)
+                # (patient, event)
+                mask = torch.sum(x, dim=2) != 0
+
+            # for case 2: [[code1, code2], [code3, ...], ...]
+            elif (dim_ == 3) and (type_ == str):
+                x = self.feat_tokenizers[feature_key].batch_encode_3d(
+                    kwargs[feature_key]
+                )
+                # (patient, visit, event)
+                x = torch.tensor(x, dtype=torch.long, device=self.device)
+                # (patient, visit, event, embedding_dim)
+                x = self.embeddings[feature_key](x)
+                # (patient, visit, embedding_dim)
+                x = torch.sum(x, dim=2)
+                # (patient, visit)
+                mask = torch.sum(x, dim=2) != 0
+
+            # for case 3: [[1.5, 2.0, 0.0], ...]
+            elif (dim_ == 2) and (type_ in [float, int]):
+                x, mask = self.padding2d(kwargs[feature_key])
+                # (patient, event, values)
+                x = torch.tensor(x, dtype=torch.float, device=self.device)
+                # (patient, event, embedding_dim)
+                if self.use_embedding[idx]:
+                    x = self.linear_layers[feature_key](x)
+                # (patient, event)
+                mask = mask.bool().to(self.device)
+
+            # for case 4: [[[1.5, 2.0, 0.0], [1.8, 2.4, 6.0]], ...]
+            elif (dim_ == 3) and (type_ in [float, int]):
+                x, mask = self.padding3d(kwargs[feature_key])
+                # (patient, visit, event, values)
+                x = torch.tensor(x, dtype=torch.float, device=self.device)
+                # (patient, visit, embedding_dim)
+                x = torch.sum(x, dim=2)
+
+                if self.use_embedding[idx]:
+                    x = self.linear_layers[feature_key](x)
+                # (patient, event)
+                mask = mask[:, :, 0]
+                mask = mask.bool().to(self.device)
+            else:
+                raise NotImplementedError
+
+            if self.static_dim > 0:
+                static = torch.tensor(
+                    kwargs[self.static_key], dtype=torch.float, device=self.device
+                )
+                x, decov = self.concare[feature_key](x, static=static, mask=mask)
+            else:
+                x, decov = self.concare[feature_key](x, mask=mask)
+            patient_emb.append(x)
+            decov_loss += decov
+
+        patient_emb = torch.cat(patient_emb, dim=1)
+        # (patient, label_size)
+        logits = self.fc(patient_emb)
+        # obtain y_true, loss, y_prob
+        y_true = self.prepare_labels(kwargs[self.label_key], self.label_tokenizer)
+        loss_task = self.get_loss_function()(logits, y_true)
+        loss = decov_loss + loss_task
+        y_prob = self.prepare_y_prob(logits)
+        # return {
+        #     "loss": loss,
+        #     "loss_task": loss_task,
+        #     "decov_loss": decov_loss,
+        #     "y_prob": y_prob,
+        #     "y_true": y_true,
+        # }
+        return {
+            "loss": loss,
+            "y_prob": y_prob,
+            "y_true": y_true,
+        }
+
+
+if __name__ == "__main__":
+    from pyhealth.datasets import SampleEHRDataset
+
+    samples = [
+        {
+            "patient_id": "patient-0",
+            "visit_id": "visit-0",
+            # "single_vector": [1, 2, 3],
+            "list_codes": ["505800458", "50580045810", "50580045811"],  # NDC
+            "list_vectors": [[1.0, 2.55, 3.4], [4.1, 5.5, 6.0]],
+            "list_list_codes": [["A05B", "A05C", "A06A"], ["A11D", "A11E"]],  # ATC-4
+            "list_list_vectors": [
+                [[1.8, 2.25, 3.41], [4.50, 5.9, 6.0]],
+                [[7.7, 8.5, 9.4]],
+            ],
+            "label": 1,
+            "demographic": [1.0, 2.0, 1.3],
+        },
+        {
+            "patient_id": "patient-0",
+            "visit_id": "visit-1",
+            # "single_vector": [1, 5, 8],
+            "list_codes": [
+                "55154191800",
+                "551541928",
+                "55154192800",
+                "705182798",
+                "70518279800",
+            ],
+            "list_vectors": [[1.4, 3.2, 3.5], [4.1, 5.9, 1.7], [4.5, 5.9, 1.7]],
+            "list_list_codes": [["A04A", "B035", "C129"]],
+            "list_list_vectors": [
+                [[1.0, 2.8, 3.3], [4.9, 5.0, 6.6], [7.7, 8.4, 1.3], [7.7, 8.4, 1.3]],
+            ],
+            "label": 0,
+            "demographic": [1.0, 2.0, 1.3],
+        },
+    ]
+
+    # dataset
+    dataset = SampleEHRDataset(samples=samples, dataset_name="test")
+
+    # data loader
+    from pyhealth.datasets import get_dataloader
+
+    train_loader = get_dataloader(dataset, batch_size=2, shuffle=True)
+
+    # model
+    model = ConCare(
+        dataset=dataset,
+        feature_keys=[
+            "list_codes",
+            "list_vectors",
+            "list_list_codes",
+            # "list_list_vectors",
+        ],
+        static_key="demographic",
+        label_key="label",
+        use_embedding=[True, False, True],
+        mode="binary",
+        hidden_dim=64,
+    )
+
+    # data batch
+    data_batch = next(iter(train_loader))
+
+    # try the model
+    ret = model(**data_batch)
+    print(ret)
+
+    # try loss backward
+    ret["loss"].backward()
diff --git a/pyhealth/models/grasp.py b/pyhealth/models/grasp.py
new file mode 100644
index 00000000..bcbc1aee
--- /dev/null
+++ b/pyhealth/models/grasp.py
@@ -0,0 +1,655 @@
+from typing import List, Tuple, Dict, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.utils.rnn as rnn_utils
+
+from pyhealth.datasets import SampleEHRDataset
+from pyhealth.models import BaseModel
+from pyhealth.models.utils import get_last_visit
+from pyhealth.models import ConCareLayer, RNNLayer
+
+from sklearn.neighbors import kneighbors_graph
+import math
+import copy
+import random
+import numpy as np
+
+
+def random_init(dataset, num_centers, device):
+    num_points = dataset.size(0)
+    dimension = dataset.size(1)
+    # print("random size", dataset.size())
+    # print("numcenter", num_centers)
+
+    indices = torch.tensor(
+        np.array(random.sample(range(num_points), k=num_centers)), dtype=torch.long
+    )
+
+    centers = torch.gather(
+        dataset, 0, indices.view(-1, 1).expand(-1, dimension).to(device=device)
+    )
+    return centers
+
+
+# Compute for each data point the closest center
+def compute_codes(dataset, centers):
+    num_points = dataset.size(0)
+    dimension = dataset.size(1)
+    num_centers = centers.size(0)
+
+    # print("size:", dataset.size(), centers.size())
+    # 5e8 should vary depending on the free memory on the GPU
+    # Ideally, automatically ;)
+    chunk_size = int(5e8 / num_centers)
+    codes = torch.zeros(num_points, dtype=torch.long)
+    centers_t = torch.transpose(centers, 0, 1)
+    centers_norms = torch.sum(centers**2, dim=1).view(1, -1)
+    for i in range(0, num_points, chunk_size):
+        begin = i
+        end = min(begin + chunk_size, num_points)
+        dataset_piece = dataset[begin:end, :]
+        dataset_norms = torch.sum(dataset_piece**2, dim=1).view(-1, 1)
+        distances = torch.mm(dataset_piece, centers_t)
+        distances *= -2.0
+        distances += dataset_norms
+        distances += centers_norms
+        _, min_ind = torch.min(distances, dim=1)
+        codes[begin:end] = min_ind
+    return codes
+
+
+# Compute new centers as means of the data points forming the clusters
+def update_centers(dataset, codes, num_centers, device):
+    num_points = dataset.size(0)
+    dimension = dataset.size(1)
+    centers = torch.zeros(num_centers, dimension, dtype=torch.float).to(device=device)
+    cnt = torch.zeros(num_centers, dtype=torch.float)
+    centers.scatter_add_(
+        0, codes.view(-1, 1).expand(-1, dimension).to(device=device), dataset
+    )
+    cnt.scatter_add_(0, codes, torch.ones(num_points, dtype=torch.float))
+    # Avoiding division by zero
+    # Not necessary if there are no duplicates among the data points
+    cnt = torch.where(cnt > 0.5, cnt, torch.ones(num_centers, dtype=torch.float))
+    centers /= cnt.view(-1, 1).to(device=device)
+    return centers
+
+
+def cluster(dataset, num_centers, device):
+    centers = random_init(dataset, num_centers, device)
+    codes = compute_codes(dataset, centers)
+    num_iterations = 0
+    while True:
+        centers = update_centers(dataset, codes, num_centers, device)
+        new_codes = compute_codes(dataset, centers)
+        # Waiting until the clustering stops updating altogether
+        # This is too strict in practice
+        if torch.equal(codes, new_codes):
+            break
+        if num_iterations > 1000:
+            break
+        codes = new_codes
+    return centers, codes
+
+
+class GraphConvolution(nn.Module):
+    def __init__(self, in_features, out_features, bias=True):
+        super(GraphConvolution, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = nn.Parameter(torch.Tensor(in_features, out_features).float())
+        if bias:
+            self.bias = nn.Parameter(torch.Tensor(out_features).float())
+        else:
+            self.register_parameter("bias", None)
+        self.initialize_parameters()
+
+    def initialize_parameters(self):
+        std = 1.0 / math.sqrt(self.weight.size(1))
+        self.weight.data.uniform_(-std, std)
+        if self.bias is not None:
+            self.bias.data.uniform_(-std, std)
+
+    def forward(self, adj, x, device):
+        y = torch.mm(x.float(), self.weight.float())
+        output = torch.mm(adj.float(), y.float())
+        if self.bias is not None:
+            return output + self.bias.float().to(device=device)
+        else:
+            return output
+
+
+def clones(module, N):
+    "Produce N identical layers."
+    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
+
+
+class GRASPLayer(nn.Module):
+    """GRASPLayer layer.
+
+    Paper: Liantao Ma et al. GRASP: generic framework for health status representation learning based on incorporating knowledge from similar patients. AAAI 2021.
+
+    This layer is used in the GRASP model. But it can also be used as a
+    standalone layer.
+
+    Args:
+        input_dim: dynamic feature size.
+        static_dim: static feature size, if 0, then no static feature is used.
+        hidden_dim: hidden dimension of the GRASP layer, default 128.
+        cluster_num: number of clusters, default 12. The cluster_num should be no more than the number of samples.
+        dropout: dropout rate, default 0.5.
+        block: the backbone model used in the GRASP layer ('ConCare', 'LSTM' or 'GRU'), default 'ConCare'.
+
+    Examples:
+        >>> from pyhealth.models import GRASPLayer
+        >>> input = torch.randn(3, 128, 64)  # [batch size, sequence len, feature_size]
+        >>> layer = GRASPLayer(64, cluster_num=2)
+        >>> c = layer(input)
+        >>> c.shape
+        torch.Size([3, 128])
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        static_dim: int = 0,
+        hidden_dim: int = 128,
+        cluster_num: int = 2,
+        dropout: int = 0.5,
+        block: int = "ConCare",
+    ):
+        super(GRASPLayer, self).__init__()
+
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.cluster_num = cluster_num
+        self.dropout = dropout
+        self.block = block
+
+        if self.block == "ConCare":
+            self.backbone = ConCareLayer(
+                input_dim, static_dim, hidden_dim, hidden_dim, dropout=0
+            )
+        elif self.block == "GRU":
+            self.backbone = RNNLayer(input_dim, hidden_dim, rnn_type="GRU", dropout=0)
+        elif self.block == "LSTM":
+            self.backbone = RNNLayer(input_dim, hidden_dim, rnn_type="LSTM", dropout=0)
+
+        self.relu = nn.ReLU()
+        self.tanh = nn.Tanh()
+        self.sigmoid = nn.Sigmoid()
+        self.dropout = nn.Dropout(dropout)
+        self.weight1 = nn.Linear(self.hidden_dim, 1)
+        self.weight2 = nn.Linear(self.hidden_dim, 1)
+        self.GCN = GraphConvolution(self.hidden_dim, self.hidden_dim, bias=True)
+        self.GCN.initialize_parameters()
+        self.GCN_2 = GraphConvolution(self.hidden_dim, self.hidden_dim, bias=True)
+        self.GCN_2.initialize_parameters()
+        self.A_mat = None
+
+        self.bn = nn.BatchNorm1d(self.hidden_dim)
+
+    def sample_gumbel(self, shape, eps=1e-20):
+        U = torch.rand(shape)
+
+        return -torch.log(-torch.log(U + eps) + eps)
+
+    def gumbel_softmax_sample(self, logits, temperature, device):
+        y = logits + self.sample_gumbel(logits.size()).to(device=device)
+        return torch.softmax(y / temperature, dim=-1)
+
+    def gumbel_softmax(self, logits, temperature, device, hard=False):
+        """
+        ST-gumple-softmax
+        input: [*, n_class]
+        return: flatten --> [*, n_class] an one-hot vector
+        """
+        y = self.gumbel_softmax_sample(logits, temperature, device)
+
+        if not hard:
+            return y.view(-1, self.cluster_num)
+
+        shape = y.size()
+        _, ind = y.max(dim=-1)
+        y_hard = torch.zeros_like(y).view(-1, shape[-1])
+        y_hard.scatter_(1, ind.view(-1, 1), 1)
+        y_hard = y_hard.view(*shape)
+        # Set gradients w.r.t. y_hard gradients w.r.t. y
+        y_hard = (y_hard - y).detach() + y
+        return y_hard
+
+    def grasp_encoder(self, input, static=None, mask=None):
+
+        if self.block == "ConCare":
+            hidden_t, _ = self.backbone(input, mask=mask, static=static)
+        else:
+            hidden_t, _ = self.backbone(input, mask)
+        hidden_t = torch.squeeze(hidden_t, 0)
+
+        centers, codes = cluster(hidden_t, self.cluster_num, input.device)
+
+        if self.A_mat == None:
+            A_mat = np.eye(self.cluster_num)
+        else:
+            A_mat = kneighbors_graph(
+                np.array(centers.detach().cpu().numpy()),
+                20,
+                mode="connectivity",
+                include_self=False,
+            ).toarray()
+
+        adj_mat = torch.tensor(A_mat).to(device=input.device)
+
+        e = self.relu(torch.matmul(hidden_t, centers.transpose(0, 1)))  # b clu_num
+
+        scores = self.gumbel_softmax(e, temperature=1, device=input.device, hard=True)
+        digits = torch.argmax(scores, dim=-1)  #  b
+
+        h_prime = self.relu(self.GCN(adj_mat, centers, input.device))
+        h_prime = self.relu(self.GCN_2(adj_mat, h_prime, input.device))
+
+        clu_appendix = torch.matmul(scores, h_prime)
+
+        weight1 = torch.sigmoid(self.weight1(clu_appendix))
+        weight2 = torch.sigmoid(self.weight2(hidden_t))
+        weight1 = weight1 / (weight1 + weight2)
+        weight2 = 1 - weight1
+
+        final_h = weight1 * clu_appendix + weight2 * hidden_t
+        out = final_h
+        return out
+
+    def forward(
+        self,
+        x: torch.tensor,
+        static: Optional[torch.tensor] = None,
+        mask: Optional[torch.tensor] = None,
+    ) -> Tuple[torch.tensor]:
+        """Forward propagation.
+
+        Args:
+            x: a tensor of shape [batch size, sequence len, input_dim].
+            static: a tensor of shape [batch size, static_dim].
+            mask: an optional tensor of shape [batch size, sequence len], where
+                1 indicates valid and 0 indicates invalid.
+
+        Returns:
+            output: a tensor of shape [batch size, fusion_dim] representing the
+                patient embedding.
+        """
+        # rnn will only apply dropout between layers
+        out = self.grasp_encoder(x, static, mask)
+        out = self.dropout(out)
+        return out
+
+
+class GRASP(BaseModel):
+    """GRASP model.
+
+    Paper: Liantao Ma et al. GRASP: generic framework for health status representation learning based on incorporating knowledge from similar patients. AAAI 2021.
+
+    Note:
+        We use separate GRASP layers for different feature_keys.
+        Currently, we automatically support different input formats:
+            - code based input (need to use the embedding table later)
+            - float/int based value input
+        We follow the current convention for the GRASP model:
+            - case 1. [code1, code2, code3, ...]
+                - we will assume the code follows the order; our model will encode
+                each code into a vector and apply GRASP on the code level
+            - case 2. [[code1, code2]] or [[code1, code2], [code3, code4, code5], ...]
+                - we will assume the inner bracket follows the order; our model first
+                use the embedding table to encode each code into a vector and then use
+                average/mean pooling to get one vector for one inner bracket; then use
+                GRASP one the braket level
+            - case 3. [[1.5, 2.0, 0.0]] or [[1.5, 2.0, 0.0], [8, 1.2, 4.5], ...]
+                - this case only makes sense when each inner bracket has the same length;
+                we assume each dimension has the same meaning; we run GRASP directly
+                on the inner bracket level, similar to case 1 after embedding table
+            - case 4. [[[1.5, 2.0, 0.0]]] or [[[1.5, 2.0, 0.0], [8, 1.2, 4.5]], ...]
+                - this case only makes sense when each inner bracket has the same length;
+                we assume each dimension has the same meaning; we run GRASP directly
+                on the inner bracket level, similar to case 2 after embedding table
+
+    Args:
+        dataset: the dataset to train the model. It is used to query certain
+            information such as the set of all tokens.
+        feature_keys:  list of keys in samples to use as features,
+            e.g. ["conditions", "procedures"].
+        label_key: key in samples to use as label (e.g., "drugs").
+        mode: one of "binary", "multiclass", or "multilabel".
+        static_keys: the key in samples to use as static features, e.g. "demographics". Default is None.
+                     we only support numerical static features.
+        use_embedding: list of bools indicating whether to use embedding for each feature type,
+            e.g. [True, False].
+        embedding_dim: the embedding dimension. Default is 128.
+        hidden_dim: the hidden dimension of the GRASP layer. Default is 128.
+        cluster_num: the number of clusters. Default is 10. Note that batch size should be greater than cluster_num.
+        **kwargs: other parameters for the GRASP layer.
+
+
+    Examples:
+        >>> from pyhealth.datasets import SampleEHRDataset
+        >>> samples = [
+        ...         {
+        ...             "patient_id": "patient-0",
+        ...             "visit_id": "visit-0",
+        ...             "list_codes": ["505800458", "50580045810", "50580045811"],  # NDC
+        ...             "list_vectors": [[1.0, 2.55, 3.4], [4.1, 5.5, 6.0]],
+        ...             "list_list_codes": [["A05B", "A05C", "A06A"], ["A11D", "A11E"]],  # ATC-4
+        ...             "list_list_vectors": [
+        ...                 [[1.8, 2.25, 3.41], [4.50, 5.9, 6.0]],
+        ...                 [[7.7, 8.5, 9.4]],
+        ...             ],
+        ...             "demographic": [0.0, 2.0, 1.5],
+        ...             "label": 1,
+        ...         },
+        ...         {
+        ...             "patient_id": "patient-0",
+        ...             "visit_id": "visit-1",
+        ...             "list_codes": [
+        ...                 "55154191800",
+        ...                 "551541928",
+        ...                 "55154192800",
+        ...                 "705182798",
+        ...                 "70518279800",
+        ...             ],
+        ...             "list_vectors": [[1.4, 3.2, 3.5], [4.1, 5.9, 1.7], [4.5, 5.9, 1.7]],
+        ...             "list_list_codes": [["A04A", "B035", "C129"]],
+        ...             "list_list_vectors": [
+        ...                 [[1.0, 2.8, 3.3], [4.9, 5.0, 6.6], [7.7, 8.4, 1.3], [7.7, 8.4, 1.3]],
+        ...             ],
+        ...             "demographic": [0.0, 2.0, 1.5],
+        ...             "label": 0,
+        ...         },
+        ...     ]
+        >>> dataset = SampleEHRDataset(samples=samples, dataset_name="test")
+        >>>
+        >>> from pyhealth.models import GRASP
+        >>> model = GRASP(
+        ...         dataset=dataset,
+        ...         feature_keys=[
+        ...             "list_codes",
+        ...             "list_vectors",
+        ...             "list_list_codes",
+        ...             "list_list_vectors",
+        ...         ],
+        ...         label_key="label",
+        ...         static_key="demographic",
+        ...         use_embedding=[True, False, True, False],
+        ...         mode="binary"
+        ...     )
+        >>>
+        >>> from pyhealth.datasets import get_dataloader
+        >>> train_loader = get_dataloader(dataset, batch_size=2, shuffle=True)
+        >>> data_batch = next(iter(train_loader))
+        >>>
+        >>> ret = model(**data_batch)
+        >>> print(ret)
+        {'loss': tensor(0.7015, grad_fn=<BinaryCrossEntropyWithLogitsBackward0>), 'y_prob': tensor([[0.5124],
+        [0.5042]], grad_fn=<SigmoidBackward0>), 'y_true': tensor([[0.],
+        [1.]])}
+        >>>
+
+    """
+
+    def __init__(
+        self,
+        dataset: SampleEHRDataset,
+        feature_keys: List[str],
+        label_key: str,
+        mode: str,
+        use_embedding: List[bool],
+        static_key: Optional[str] = None,
+        embedding_dim: int = 128,
+        hidden_dim: int = 128,
+        cluster_num: int = 10,
+        **kwargs,
+    ):
+        super(GRASP, self).__init__(
+            dataset=dataset,
+            feature_keys=feature_keys,
+            label_key=label_key,
+            mode=mode,
+        )
+        self.embedding_dim = embedding_dim
+        self.hidden_dim = hidden_dim
+        self.use_embedding = use_embedding
+
+        # validate kwargs for GRASP layer
+        if "feature_size" in kwargs:
+            raise ValueError("feature_size is determined by embedding_dim")
+        if len(dataset) < 12 and "cluster_num" not in kwargs:
+            raise ValueError("cluster_num is required for small dataset, default 12")
+        if "cluster_num" in kwargs and kwargs["cluster_num"] > len(dataset):
+            raise ValueError("cluster_num must be no larger than dataset size")
+
+        # the key of self.feat_tokenizers only contains the code based inputs
+        self.feat_tokenizers = {}
+        self.static_key = static_key
+        self.label_tokenizer = self.get_label_tokenizer()
+        # the key of self.embeddings only contains the code based inputs
+        self.embeddings = nn.ModuleDict()
+        # the key of self.linear_layers only contains the float/int based inputs
+        self.linear_layers = nn.ModuleDict()
+
+        self.static_dim = 0
+        if self.static_key is not None:
+            self.static_dim = self.dataset.input_info[self.static_key]["len"]
+
+        self.grasp = nn.ModuleDict()
+        # add feature GRASP layers
+        for idx, feature_key in enumerate(self.feature_keys):
+            input_info = self.dataset.input_info[feature_key]
+            # sanity check
+            if input_info["type"] not in [str, float, int]:
+                raise ValueError(
+                    "GRASP only supports str code, float and int as input types"
+                )
+            elif (input_info["type"] == str) and (input_info["dim"] not in [2, 3]):
+                raise ValueError(
+                    "GRASP only supports 2-dim or 3-dim str code as input types"
+                )
+            elif (input_info["type"] == str) and (use_embedding[idx] == False):
+                raise ValueError(
+                    "GRASP only supports embedding for str code as input types"
+                )
+            elif (input_info["type"] in [float, int]) and (
+                input_info["dim"] not in [2, 3]
+            ):
+                raise ValueError(
+                    "GRASP only supports 2-dim or 3-dim float and int as input types"
+                )
+
+            # for code based input, we need Type
+            # for float/int based input, we need Type, input_dim
+            if use_embedding[idx]:
+                self.add_feature_transform_layer(feature_key, input_info)
+                self.grasp[feature_key] = GRASPLayer(
+                    input_dim=embedding_dim,
+                    static_dim=self.static_dim,
+                    hidden_dim=self.hidden_dim,
+                    cluster_num=cluster_num,
+                    **kwargs,
+                )
+            else:
+                self.grasp[feature_key] = GRASPLayer(
+                    input_dim=input_info["len"],
+                    static_dim=self.static_dim,
+                    hidden_dim=self.hidden_dim,
+                    cluster_num=cluster_num,
+                    **kwargs,
+                )
+
+        output_size = self.get_output_size(self.label_tokenizer)
+        self.fc = nn.Linear(len(self.feature_keys) * self.hidden_dim, output_size)
+
+    def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
+        """Forward propagation.
+
+        The label `kwargs[self.label_key]` is a list of labels for each patient.
+
+        Args:
+            **kwargs: keyword arguments for the model. The keys must contain
+                all the feature keys and the label key.
+
+        Returns:
+            A dictionary with the following keys:
+                loss: a scalar tensor representing the final loss.
+                y_prob: a tensor representing the predicted probabilities.
+                y_true: a tensor representing the true labels.
+        """
+        patient_emb = []
+        for idx, feature_key in enumerate(self.feature_keys):
+            input_info = self.dataset.input_info[feature_key]
+            dim_, type_ = input_info["dim"], input_info["type"]
+
+            # for case 1: [code1, code2, code3, ...]
+            if (dim_ == 2) and (type_ == str):
+                x = self.feat_tokenizers[feature_key].batch_encode_2d(
+                    kwargs[feature_key]
+                )
+                # (patient, event)
+                x = torch.tensor(x, dtype=torch.long, device=self.device)
+                # (patient, event, embedding_dim)
+                x = self.embeddings[feature_key](x)
+                # (patient, event)
+                mask = torch.sum(x, dim=2) != 0
+
+            # for case 2: [[code1, code2], [code3, ...], ...]
+            elif (dim_ == 3) and (type_ == str):
+                x = self.feat_tokenizers[feature_key].batch_encode_3d(
+                    kwargs[feature_key]
+                )
+                # (patient, visit, event)
+                x = torch.tensor(x, dtype=torch.long, device=self.device)
+                # (patient, visit, event, embedding_dim)
+                x = self.embeddings[feature_key](x)
+                # (patient, visit, embedding_dim)
+                x = torch.sum(x, dim=2)
+                # (patient, visit)
+                mask = torch.sum(x, dim=2) != 0
+
+            # for case 3: [[1.5, 2.0, 0.0], ...]
+            elif (dim_ == 2) and (type_ in [float, int]):
+                x, mask = self.padding2d(kwargs[feature_key])
+                # (patient, event, values)
+                x = torch.tensor(x, dtype=torch.float, device=self.device)
+                # (patient, event, embedding_dim)
+                if self.use_embedding[idx]:
+                    x = self.linear_layers[feature_key](x)
+                # (patient, event)
+                mask = mask.bool().to(self.device)
+
+            # for case 4: [[[1.5, 2.0, 0.0], [1.8, 2.4, 6.0]], ...]
+            elif (dim_ == 3) and (type_ in [float, int]):
+                x, mask = self.padding3d(kwargs[feature_key])
+                # (patient, visit, event, values)
+                x = torch.tensor(x, dtype=torch.float, device=self.device)
+                # (patient, visit, embedding_dim)
+                x = torch.sum(x, dim=2)
+
+                if self.use_embedding[idx]:
+                    x = self.linear_layers[feature_key](x)
+                # (patient, event)
+                mask = mask[:, :, 0]
+                mask = mask.bool().to(self.device)
+            else:
+                raise NotImplementedError
+
+            if self.static_dim > 0:
+                static = torch.tensor(
+                    kwargs[self.static_key], dtype=torch.float, device=self.device
+                )
+                x = self.grasp[feature_key](x, static=static, mask=mask)
+            else:
+                x = self.grasp[feature_key](x, mask=mask)
+            patient_emb.append(x)
+
+        patient_emb = torch.cat(patient_emb, dim=1)
+        # (patient, label_size)
+        logits = self.fc(patient_emb)
+        # obtain y_true, loss, y_prob
+        y_true = self.prepare_labels(kwargs[self.label_key], self.label_tokenizer)
+        loss = self.get_loss_function()(logits, y_true)
+        y_prob = self.prepare_y_prob(logits)
+        return {
+            "loss": loss,
+            "y_prob": y_prob,
+            "y_true": y_true,
+        }
+
+
+if __name__ == "__main__":
+    from pyhealth.datasets import SampleEHRDataset
+
+    samples = [
+        {
+            "patient_id": "patient-0",
+            "visit_id": "visit-0",
+            # "single_vector": [1, 2, 3],
+            "list_codes": ["505800458", "50580045810", "50580045811"],  # NDC
+            "list_vectors": [[1.0, 2.55, 3.4], [4.1, 5.5, 6.0]],
+            "list_list_codes": [["A05B", "A05C", "A06A"], ["A11D", "A11E"]],  # ATC-4
+            "list_list_vectors": [
+                [[1.8, 2.25, 3.41], [4.50, 5.9, 6.0]],
+                [[7.7, 8.5, 9.4]],
+            ],
+            "label": 1,
+            "demographic": [1.0, 2.0, 1.3],
+        },
+        {
+            "patient_id": "patient-0",
+            "visit_id": "visit-1",
+            # "single_vector": [1, 5, 8],
+            "list_codes": [
+                "55154191800",
+                "551541928",
+                "55154192800",
+                "705182798",
+                "70518279800",
+            ],
+            "list_vectors": [[1.4, 3.2, 3.5], [4.1, 5.9, 1.7], [4.5, 5.9, 1.7]],
+            "list_list_codes": [["A04A", "B035", "C129"]],
+            "list_list_vectors": [
+                [[1.0, 2.8, 3.3], [4.9, 5.0, 6.6], [7.7, 8.4, 1.3], [7.7, 8.4, 1.3]],
+            ],
+            "label": 0,
+            "demographic": [1.0, 2.0, 1.3],
+        },
+    ]
+
+    # dataset
+    dataset = SampleEHRDataset(samples=samples, dataset_name="test")
+
+    # data loader
+    from pyhealth.datasets import get_dataloader
+
+    train_loader = get_dataloader(dataset, batch_size=2, shuffle=True)
+
+    # model
+    model = GRASP(
+        dataset=dataset,
+        feature_keys=[
+            "list_codes",
+            "list_vectors",
+            "list_list_codes",
+            # "list_list_vectors",
+        ],
+        static_key="demographic",
+        label_key="label",
+        use_embedding=[True, False, True],
+        mode="binary",
+        cluster_num=2,
+    )
+
+    # data batch
+    data_batch = next(iter(train_loader))
+
+    # try the model
+    ret = model(**data_batch)
+    print(ret)
+
+    # try loss backward
+    ret["loss"].backward()
diff --git a/pyhealth/models/stagenet.py b/pyhealth/models/stagenet.py
new file mode 100644
index 00000000..66ec5c39
--- /dev/null
+++ b/pyhealth/models/stagenet.py
@@ -0,0 +1,613 @@
+from typing import List, Tuple, Dict, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.utils.rnn as rnn_utils
+
+from pyhealth.datasets import SampleEHRDataset
+from pyhealth.models import BaseModel
+from pyhealth.models.utils import get_last_visit
+import torch.nn.functional as F
+
+# VALID_OPERATION_LEVEL = ["visit", "event"]
+
+
+class StageNetLayer(nn.Module):
+    """StageNet layer.
+
+    Paper: Stagenet: Stage-aware neural networks for health risk prediction. WWW 2020.
+
+    This layer is used in the StageNet model. But it can also be used as a
+    standalone layer.
+
+    Args:
+        input_dim: dynamic feature size.
+        chunk_size: the chunk size for the StageNet layer. Default is 128.
+        levels: the number of levels for the StageNet layer. levels * chunk_size = hidden_dim in the RNN. Smaller chunk size and more levels can capture more detailed patient status variations. Default is 3.
+        conv_size: the size of the convolutional kernel. Default is 10.
+        dropconnect: the dropout rate for the dropconnect. Default is 0.3.
+        dropout: the dropout rate for the dropout. Default is 0.3.
+        dropres: the dropout rate for the residual connection. Default is 0.3.
+
+    Examples:
+        >>> from pyhealth.models import StageNetLayer
+        >>> input = torch.randn(3, 128, 64)  # [batch size, sequence len, feature_size]
+        >>> layer = StageNetLayer(64)
+        >>> c, _, _ = layer(input)
+        >>> c.shape
+        torch.Size([3, 384])
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        chunk_size: int = 128,
+        conv_size: int = 10,
+        levels: int = 3,
+        dropconnect: int = 0.3,
+        dropout: int = 0.3,
+        dropres: int = 0.3,
+    ):
+        super(StageNetLayer, self).__init__()
+
+        self.dropout = dropout
+        self.dropconnect = dropconnect
+        self.dropres = dropres
+        self.input_dim = input_dim
+        self.hidden_dim = chunk_size * levels
+        self.conv_dim = self.hidden_dim
+        self.conv_size = conv_size
+        # self.output_dim = output_dim
+        self.levels = levels
+        self.chunk_size = chunk_size
+
+        self.kernel = nn.Linear(
+            int(input_dim + 1), int(self.hidden_dim * 4 + levels * 2)
+        )
+        nn.init.xavier_uniform_(self.kernel.weight)
+        nn.init.zeros_(self.kernel.bias)
+        self.recurrent_kernel = nn.Linear(
+            int(self.hidden_dim + 1), int(self.hidden_dim * 4 + levels * 2)
+        )
+        nn.init.orthogonal_(self.recurrent_kernel.weight)
+        nn.init.zeros_(self.recurrent_kernel.bias)
+
+        self.nn_scale = nn.Linear(int(self.hidden_dim), int(self.hidden_dim // 6))
+        self.nn_rescale = nn.Linear(int(self.hidden_dim // 6), int(self.hidden_dim))
+        self.nn_conv = nn.Conv1d(
+            int(self.hidden_dim), int(self.conv_dim), int(conv_size), 1
+        )
+        # self.nn_output = nn.Linear(int(self.conv_dim), int(output_dim))
+
+        if self.dropconnect:
+            self.nn_dropconnect = nn.Dropout(p=dropconnect)
+            self.nn_dropconnect_r = nn.Dropout(p=dropconnect)
+        if self.dropout:
+            self.nn_dropout = nn.Dropout(p=dropout)
+            self.nn_dropres = nn.Dropout(p=dropres)
+
+    def cumax(self, x, mode="l2r"):
+        if mode == "l2r":
+            x = torch.softmax(x, dim=-1)
+            x = torch.cumsum(x, dim=-1)
+            return x
+        elif mode == "r2l":
+            x = torch.flip(x, [-1])
+            x = torch.softmax(x, dim=-1)
+            x = torch.cumsum(x, dim=-1)
+            return torch.flip(x, [-1])
+        else:
+            return x
+
+    def step(self, inputs, c_last, h_last, interval, device):
+        x_in = inputs.to(device=device)
+
+        # Integrate inter-visit time intervals
+        interval = interval.unsqueeze(-1).to(device=device)
+        x_out1 = self.kernel(torch.cat((x_in, interval), dim=-1)).to(device)
+        x_out2 = self.recurrent_kernel(
+            torch.cat((h_last.to(device=device), interval), dim=-1)
+        )
+
+        if self.dropconnect:
+            x_out1 = self.nn_dropconnect(x_out1)
+            x_out2 = self.nn_dropconnect_r(x_out2)
+        x_out = x_out1 + x_out2
+        f_master_gate = self.cumax(x_out[:, : self.levels], "l2r")
+        f_master_gate = f_master_gate.unsqueeze(2).to(device=device)
+        i_master_gate = self.cumax(x_out[:, self.levels : self.levels * 2], "r2l")
+        i_master_gate = i_master_gate.unsqueeze(2)
+        x_out = x_out[:, self.levels * 2 :]
+        x_out = x_out.reshape(-1, self.levels * 4, self.chunk_size)
+        f_gate = torch.sigmoid(x_out[:, : self.levels]).to(device=device)
+        i_gate = torch.sigmoid(x_out[:, self.levels : self.levels * 2]).to(
+            device=device
+        )
+        o_gate = torch.sigmoid(x_out[:, self.levels * 2 : self.levels * 3])
+        c_in = torch.tanh(x_out[:, self.levels * 3 :]).to(device=device)
+        c_last = c_last.reshape(-1, self.levels, self.chunk_size).to(device=device)
+        overlap = (f_master_gate * i_master_gate).to(device=device)
+        c_out = (
+            overlap * (f_gate * c_last + i_gate * c_in)
+            + (f_master_gate - overlap) * c_last
+            + (i_master_gate - overlap) * c_in
+        )
+        h_out = o_gate * torch.tanh(c_out)
+        c_out = c_out.reshape(-1, self.hidden_dim)
+        h_out = h_out.reshape(-1, self.hidden_dim)
+        out = torch.cat([h_out, f_master_gate[..., 0], i_master_gate[..., 0]], 1)
+        return out, c_out, h_out
+
+    def forward(
+        self,
+        x: torch.tensor,
+        time: Optional[torch.tensor] = None,
+        mask: Optional[torch.tensor] = None,
+    ) -> Tuple[torch.tensor]:
+        """Forward propagation.
+
+        Args:
+            x: a tensor of shape [batch size, sequence len, input_dim].
+            static: a tensor of shape [batch size, static_dim].
+            mask: an optional tensor of shape [batch size, sequence len], where
+                1 indicates valid and 0 indicates invalid.
+
+        Returns:
+            last_output: a tensor of shape [batch size, chunk_size*levels] representing the
+                patient embedding.
+            outputs: a tensor of shape [batch size, sequence len, chunk_size*levels] representing the patient at each time step.
+        """
+        # rnn will only apply dropout between layers
+        batch_size, time_step, feature_dim = x.size()
+        device = x.device
+        if time == None:
+            time = torch.ones(batch_size, time_step)
+        time = time.reshape(batch_size, time_step)
+        c_out = torch.zeros(batch_size, self.hidden_dim)
+        h_out = torch.zeros(batch_size, self.hidden_dim)
+
+        tmp_h = (
+            torch.zeros_like(h_out, dtype=torch.float32)
+            .view(-1)
+            .repeat(self.conv_size)
+            .view(self.conv_size, batch_size, self.hidden_dim)
+        )
+        tmp_dis = torch.zeros((self.conv_size, batch_size))
+        h = []
+        origin_h = []
+        distance = []
+        for t in range(time_step):
+            out, c_out, h_out = self.step(x[:, t, :], c_out, h_out, time[:, t], device)
+            cur_distance = 1 - torch.mean(
+                out[..., self.hidden_dim : self.hidden_dim + self.levels], -1
+            )
+            origin_h.append(out[..., : self.hidden_dim])
+            tmp_h = torch.cat(
+                (
+                    tmp_h[1:].to(device=device),
+                    out[..., : self.hidden_dim].unsqueeze(0).to(device=device),
+                ),
+                0,
+            )
+            tmp_dis = torch.cat(
+                (
+                    tmp_dis[1:].to(device=device),
+                    cur_distance.unsqueeze(0).to(device=device),
+                ),
+                0,
+            )
+            distance.append(cur_distance)
+
+            # Re-weighted convolution operation
+            local_dis = tmp_dis.permute(1, 0)
+            local_dis = torch.cumsum(local_dis, dim=1)
+            local_dis = torch.softmax(local_dis, dim=1)
+            local_h = tmp_h.permute(1, 2, 0)
+            local_h = local_h * local_dis.unsqueeze(1)
+
+            # Re-calibrate Progression patterns
+            local_theme = torch.mean(local_h, dim=-1)
+            local_theme = self.nn_scale(local_theme).to(device)
+            local_theme = torch.relu(local_theme)
+            local_theme = self.nn_rescale(local_theme).to(device)
+            local_theme = torch.sigmoid(local_theme)
+
+            local_h = self.nn_conv(local_h).squeeze(-1)
+            local_h = local_theme * local_h
+            h.append(local_h)
+
+        origin_h = torch.stack(origin_h).permute(1, 0, 2)
+        rnn_outputs = torch.stack(h).permute(1, 0, 2)
+        if self.dropres > 0.0:
+            origin_h = self.nn_dropres(origin_h)
+        rnn_outputs = rnn_outputs + origin_h
+        rnn_outputs = rnn_outputs.contiguous().view(-1, rnn_outputs.size(-1))
+        if self.dropout > 0.0:
+            rnn_outputs = self.nn_dropout(rnn_outputs)
+
+        output = rnn_outputs.contiguous().view(batch_size, time_step, self.hidden_dim)
+        last_output = get_last_visit(output, mask)
+
+        return last_output, output, torch.stack(distance)
+
+
+class StageNet(BaseModel):
+    """StageNet model.
+
+    Paper: Junyi Gao et al. Stagenet: Stage-aware neural networks for health risk prediction. WWW 2020.
+
+    Note:
+        We use separate StageNet layers for different feature_keys.
+        Currently, we automatically support different input formats:
+            - code based input (need to use the embedding table later)
+            - float/int based value input
+        We follow the current convention for the StageNet model:
+            - case 1. [code1, code2, code3, ...]
+                - we will assume the code follows the order; our model will encode
+                each code into a vector and apply StageNet on the code level
+            - case 2. [[code1, code2]] or [[code1, code2], [code3, code4, code5], ...]
+                - we will assume the inner bracket follows the order; our model first
+                use the embedding table to encode each code into a vector and then use
+                average/mean pooling to get one vector for one inner bracket; then use
+                StageNet one the braket level
+            - case 3. [[1.5, 2.0, 0.0]] or [[1.5, 2.0, 0.0], [8, 1.2, 4.5], ...]
+                - this case only makes sense when each inner bracket has the same length;
+                we assume each dimension has the same meaning; we run StageNet directly
+                on the inner bracket level, similar to case 1 after embedding table
+            - case 4. [[[1.5, 2.0, 0.0]]] or [[[1.5, 2.0, 0.0], [8, 1.2, 4.5]], ...]
+                - this case only makes sense when each inner bracket has the same length;
+                we assume each dimension has the same meaning; we run StageNet directly
+                on the inner bracket level, similar to case 2 after embedding table
+        The time interval information specified by time_keys will be used to calculate the memory decay between each visit. If time_keys is None, all visits are treated as the same time interval. For each feature, the time interval should be a two-dimensional float array with shape (time_step, 1).
+
+    Args:
+        dataset: the dataset to train the model. It is used to query certain
+            information such as the set of all tokens.
+        feature_keys:  list of keys in samples to use as features,
+            e.g. ["conditions", "procedures"].
+        label_key: key in samples to use as label (e.g., "drugs").
+        mode: one of "binary", "multiclass", or "multilabel".
+        time_keys: list of keys in samples to use as time interval information for each feature, Default is None. If none, all visits are treated as the same time interval.
+        embedding_dim: the embedding dimension. Default is 128.
+        chunk_size: the chunk size for the StageNet layer. Default is 128.
+        levels: the number of levels for the StageNet layer. levels * chunk_size = hidden_dim in the RNN. Smaller chunk size and more levels can capture more detailed patient status variations. Default is 3.
+        **kwargs: other parameters for the StageNet layer.
+
+
+    Examples:
+        >>> from pyhealth.datasets import SampleEHRDataset
+        >>> samples = [
+        ...     {
+        ...         "patient_id": "patient-0",
+        ...         "visit_id": "visit-0",
+        ...         # "single_vector": [1, 2, 3],
+        ...         "list_codes": ["505800458", "50580045810", "50580045811"],  # NDC
+        ...         "list_vectors": [[1.0, 2.55, 3.4], [4.1, 5.5, 6.0]],
+        ...         "list_list_codes": [["A05B", "A05C", "A06A"], ["A11D", "A11E"]],  # ATC-4
+        ...         "list_list_vectors": [
+        ...             [[1.8, 2.25, 3.41], [4.50, 5.9, 6.0]],
+        ...             [[7.7, 8.5, 9.4]],
+        ...         ],
+        ...         "label": 1,
+        ...         "list_vectors_time": [[0.0], [1.3]],
+        ...         "list_codes_time": [[0.0], [2.0], [1.3]],
+        ...         "list_list_codes_time": [[0.0], [1.5]],
+        ...     },
+        ...     {
+        ...         "patient_id": "patient-0",
+        ...         "visit_id": "visit-1",
+        ...         # "single_vector": [1, 5, 8],
+        ...         "list_codes": [
+        ...             "55154191800",
+        ...             "551541928",
+        ...             "55154192800",
+        ...             "705182798",
+        ...             "70518279800",
+        ...         ],
+        ...         "list_vectors": [[1.4, 3.2, 3.5], [4.1, 5.9, 1.7], [4.5, 5.9, 1.7]],
+        ...         "list_list_codes": [["A04A", "B035", "C129"]],
+        ...         "list_list_vectors": [
+        ...             [[1.0, 2.8, 3.3], [4.9, 5.0, 6.6], [7.7, 8.4, 1.3], [7.7, 8.4, 1.3]],
+        ...         ],
+        ...         "label": 0,
+        ...         "list_vectors_time": [[0.0], [2.0], [1.0]],
+        ...         "list_codes_time": [[0.0], [2.0], [1.3], [1.0], [2.0]],
+        ...         "list_list_codes_time": [[0.0]],
+        ...     },
+        ... ]
+        >>>
+        >>> # dataset
+        >>> dataset = SampleEHRDataset(samples=samples, dataset_name="test")
+        >>>
+        >>> # data loader
+        >>> from pyhealth.datasets import get_dataloader
+        >>>
+        >>> train_loader = get_dataloader(dataset, batch_size=2, shuffle=True)
+        >>>
+        >>> # model
+        >>> model = StageNet(
+        ...     dataset=dataset,
+        ...     feature_keys=[
+        ...         "list_codes",
+        ...         "list_vectors",
+        ...         "list_list_codes",
+        ...         # "list_list_vectors",
+        ...     ],
+        ...     time_keys=["list_codes_time", "list_vectors_time", "list_list_codes_time"],
+        ...     label_key="label",
+        ...     mode="binary",
+        ... )
+        >>>
+        >>> from pyhealth.datasets import get_dataloader
+        >>> train_loader = get_dataloader(dataset, batch_size=2, shuffle=True)
+        >>> data_batch = next(iter(train_loader))
+        >>>
+        >>> ret = model(**data_batch)
+        >>> print(ret)
+        {'loss': tensor(0.7272, grad_fn=<BinaryCrossEntropyWithLogitsBackward0>), 'distance':..., 'y_prob': tensor([[0.4328],
+        [0.4605]], grad_fn=<SigmoidBackward0>), 'y_true': tensor([[1.],
+        [0.]])}
+        >>>
+
+    """
+
+    def __init__(
+        self,
+        dataset: SampleEHRDataset,
+        feature_keys: List[str],
+        label_key: str,
+        mode: str,
+        time_keys: List[str] = None,
+        embedding_dim: int = 128,
+        chunk_size: int = 128,
+        levels: int = 3,
+        **kwargs,
+    ):
+        super(StageNet, self).__init__(
+            dataset=dataset,
+            feature_keys=feature_keys,
+            label_key=label_key,
+            mode=mode,
+        )
+        self.embedding_dim = embedding_dim
+        self.chunk_size = chunk_size
+        self.levels = levels
+
+        # validate kwargs for StageNet layer
+        if "feature_size" in kwargs:
+            raise ValueError("feature_size is determined by embedding_dim")
+        if time_keys is not None:
+            if len(time_keys) != len(feature_keys):
+                raise ValueError(
+                    "time_keys should have the same length as feature_keys"
+                )
+
+        # the key of self.feat_tokenizers only contains the code based inputs
+        self.feat_tokenizers = {}
+        self.time_keys = time_keys
+        self.label_tokenizer = self.get_label_tokenizer()
+        # the key of self.embeddings only contains the code based inputs
+        self.embeddings = nn.ModuleDict()
+        # the key of self.linear_layers only contains the float/int based inputs
+        self.linear_layers = nn.ModuleDict()
+
+        self.stagenet = nn.ModuleDict()
+        # add feature StageNet layers
+        for feature_key in self.feature_keys:
+            input_info = self.dataset.input_info[feature_key]
+            # sanity check
+            if input_info["type"] not in [str, float, int]:
+                raise ValueError(
+                    "StageNet only supports str code, float and int as input types"
+                )
+            elif (input_info["type"] == str) and (input_info["dim"] not in [2, 3]):
+                raise ValueError(
+                    "StageNet only supports 2-dim or 3-dim str code as input types"
+                )
+            elif (input_info["type"] in [float, int]) and (
+                input_info["dim"] not in [2, 3]
+            ):
+                raise ValueError(
+                    "StageNet only supports 2-dim or 3-dim float and int as input types"
+                )
+
+            # for code based input, we need Type
+            # for float/int based input, we need Type, input_dim
+            self.add_feature_transform_layer(feature_key, input_info)
+            self.stagenet[feature_key] = StageNetLayer(
+                input_dim=embedding_dim,
+                chunk_size=self.chunk_size,
+                levels=self.levels,
+                **kwargs,
+            )
+
+        output_size = self.get_output_size(self.label_tokenizer)
+        self.fc = nn.Linear(
+            len(self.feature_keys) * self.chunk_size * self.levels, output_size
+        )
+
+    def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
+        """Forward propagation.
+
+        The label `kwargs[self.label_key]` is a list of labels for each patient.
+
+        Args:
+            **kwargs: keyword arguments for the model. The keys must contain
+                all the feature keys and the label key.
+
+        Returns:
+            A dictionary with the following keys:
+                loss: a scalar tensor representing the final loss.
+                distance: list of tensors representing the stage variation of the patient.
+                y_prob: a tensor representing the predicted probabilities.
+                y_true: a tensor representing the true labels.
+        """
+        patient_emb = []
+        distance = []
+        mask_dict = {}
+        for idx, feature_key in enumerate(self.feature_keys):
+            input_info = self.dataset.input_info[feature_key]
+            dim_, type_ = input_info["dim"], input_info["type"]
+
+            # for case 1: [code1, code2, code3, ...]
+            if (dim_ == 2) and (type_ == str):
+                x = self.feat_tokenizers[feature_key].batch_encode_2d(
+                    kwargs[feature_key]
+                )
+                # (patient, event)
+                x = torch.tensor(x, dtype=torch.long, device=self.device)
+                # (patient, event, embedding_dim)
+                x = self.embeddings[feature_key](x)
+                # (patient, event)
+                mask = torch.sum(x, dim=2) != 0
+                mask_dict[feature_key] = mask
+
+            # for case 2: [[code1, code2], [code3, ...], ...]
+            elif (dim_ == 3) and (type_ == str):
+                x = self.feat_tokenizers[feature_key].batch_encode_3d(
+                    kwargs[feature_key]
+                )
+                # (patient, visit, event)
+                x = torch.tensor(x, dtype=torch.long, device=self.device)
+                # (patient, visit, event, embedding_dim)
+                x = self.embeddings[feature_key](x)
+                # (patient, visit, embedding_dim)
+                x = torch.sum(x, dim=2)
+                # (patient, visit)
+                mask = torch.sum(x, dim=2) != 0
+                mask_dict[feature_key] = mask
+
+            # for case 3: [[1.5, 2.0, 0.0], ...]
+            elif (dim_ == 2) and (type_ in [float, int]):
+                x, mask = self.padding2d(kwargs[feature_key])
+                # (patient, event, values)
+                x = torch.tensor(x, dtype=torch.float, device=self.device)
+                # (patient, event, embedding_dim)
+                x = self.linear_layers[feature_key](x)
+                # (patient, event)
+                mask = mask.bool().to(self.device)
+                mask_dict[feature_key] = mask
+
+            # for case 4: [[[1.5, 2.0, 0.0], [1.8, 2.4, 6.0]], ...]
+            elif (dim_ == 3) and (type_ in [float, int]):
+                x, mask = self.padding3d(kwargs[feature_key])
+                # (patient, visit, event, values)
+                x = torch.tensor(x, dtype=torch.float, device=self.device)
+                # (patient, visit, embedding_dim)
+                x = torch.sum(x, dim=2)
+                x = self.linear_layers[feature_key](x)
+                # (patient, event)
+                mask = mask[:, :, 0]
+                mask = mask.bool().to(self.device)
+                mask_dict[feature_key] = mask
+
+            else:
+                raise NotImplementedError
+
+            time = None
+            if self.time_keys is not None:
+                input_info = self.dataset.input_info[self.time_keys[idx]]
+                dim_, type_ = input_info["dim"], input_info["type"]
+                if (dim_ != 2) or (type_ not in [float, int]):
+                    raise ValueError("Time interval must be 2-dim float or int.")
+                time, _ = self.padding2d(kwargs[self.time_keys[idx]])
+                time = torch.tensor(time, dtype=torch.float, device=self.device)
+            x, _, cur_dis = self.stagenet[feature_key](x, time=time, mask=mask)
+            patient_emb.append(x)
+            distance.append(cur_dis)
+
+        patient_emb = torch.cat(patient_emb, dim=1)
+        # (patient, label_size)
+        logits = self.fc(patient_emb)
+        # obtain y_true, loss, y_prob
+        y_true = self.prepare_labels(kwargs[self.label_key], self.label_tokenizer)
+        loss = self.get_loss_function()(logits, y_true)
+
+        y_prob = self.prepare_y_prob(logits)
+        # return {
+        #     "loss": loss,
+        #     "distance": distance,
+        #     "y_prob": y_prob,
+        #     "y_true": y_true,
+        # }
+        return {
+            "loss": loss,
+            "y_prob": y_prob,
+            "y_true": y_true,
+        }
+
+
+if __name__ == "__main__":
+    from pyhealth.datasets import SampleEHRDataset
+
+    samples = [
+        {
+            "patient_id": "patient-0",
+            "visit_id": "visit-0",
+            # "single_vector": [1, 2, 3],
+            "list_codes": ["505800458", "50580045810", "50580045811"],  # NDC
+            "list_vectors": [[1.0, 2.55, 3.4], [4.1, 5.5, 6.0]],
+            "list_list_codes": [["A05B", "A05C", "A06A"], ["A11D", "A11E"]],  # ATC-4
+            "list_list_vectors": [
+                [[1.8, 2.25, 3.41], [4.50, 5.9, 6.0]],
+                [[7.7, 8.5, 9.4]],
+            ],
+            "label": 1,
+            "list_vectors_time": [[0.0], [1.3]],
+            "list_codes_time": [[0.0], [2.0], [1.3]],
+            "list_list_codes_time": [[0.0], [1.5]],
+        },
+        {
+            "patient_id": "patient-0",
+            "visit_id": "visit-1",
+            # "single_vector": [1, 5, 8],
+            "list_codes": [
+                "55154191800",
+                "551541928",
+                "55154192800",
+                "705182798",
+                "70518279800",
+            ],
+            "list_vectors": [[1.4, 3.2, 3.5], [4.1, 5.9, 1.7], [4.5, 5.9, 1.7]],
+            "list_list_codes": [["A04A", "B035", "C129"]],
+            "list_list_vectors": [
+                [[1.0, 2.8, 3.3], [4.9, 5.0, 6.6], [7.7, 8.4, 1.3], [7.7, 8.4, 1.3]],
+            ],
+            "label": 0,
+            "list_vectors_time": [[0.0], [2.0], [1.0]],
+            "list_codes_time": [[0.0], [2.0], [1.3], [1.0], [2.0]],
+            "list_list_codes_time": [[0.0]],
+        },
+    ]
+
+    # dataset
+    dataset = SampleEHRDataset(samples=samples, dataset_name="test")
+
+    # data loader
+    from pyhealth.datasets import get_dataloader
+
+    train_loader = get_dataloader(dataset, batch_size=2, shuffle=True)
+
+    # model
+    model = StageNet(
+        dataset=dataset,
+        feature_keys=[
+            "list_codes",
+            "list_vectors",
+            "list_list_codes",
+            # "list_list_vectors",
+        ],
+        time_keys=["list_codes_time", "list_vectors_time", "list_list_codes_time"],
+        label_key="label",
+        mode="binary",
+    )
+
+    # data batch
+    data_batch = next(iter(train_loader))
+
+    # try the model
+    ret = model(**data_batch)
+    print(ret)
+
+    # try loss backward
+    ret["loss"].backward()
diff --git a/pyhealth/models/tcn.py b/pyhealth/models/tcn.py
new file mode 100644
index 00000000..a78f791e
--- /dev/null
+++ b/pyhealth/models/tcn.py
@@ -0,0 +1,503 @@
+from typing import List, Tuple, Dict, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.utils.rnn as rnn_utils
+
+from pyhealth.datasets import SampleEHRDataset
+from pyhealth.models import BaseModel
+from torch.nn.utils import weight_norm
+
+import numpy as np
+from pyhealth.models.utils import get_last_visit
+
+
+# VALID_OPERATION_LEVEL = ["visit", "event"]
+
+
+# From TCN original paper https://github.com/locuslab/TCN
+class Chomp1d(nn.Module):
+    def __init__(self, chomp_size):
+        super(Chomp1d, self).__init__()
+        self.chomp_size = chomp_size
+
+    def forward(self, x):
+        return x[:, :, : -self.chomp_size].contiguous()
+
+
+class TemporalBlock(nn.Module):
+    def __init__(
+        self, n_inputs, n_outputs, kernel_size, stride, dilation, padding, dropout=0.2
+    ):
+        super(TemporalBlock, self).__init__()
+        self.conv1 = weight_norm(
+            nn.Conv1d(
+                n_inputs,
+                n_outputs,
+                kernel_size,
+                stride=stride,
+                padding=padding,
+                dilation=dilation,
+            ),
+            dim=None,
+        )
+        self.chomp1 = Chomp1d(padding)
+        self.relu1 = nn.ReLU()
+        self.dropout1 = nn.Dropout(dropout)
+
+        self.conv2 = weight_norm(
+            nn.Conv1d(
+                n_outputs,
+                n_outputs,
+                kernel_size,
+                stride=stride,
+                padding=padding,
+                dilation=dilation,
+            ),
+            dim=None,
+        )
+        self.chomp2 = Chomp1d(padding)
+        self.relu2 = nn.ReLU()
+        self.dropout2 = nn.Dropout(dropout)
+
+        self.net = nn.Sequential(
+            self.conv1,
+            self.chomp1,
+            self.relu1,
+            self.dropout1,
+            self.conv2,
+            self.chomp2,
+            self.relu2,
+            self.dropout2,
+        )
+        self.downsample = (
+            nn.Conv1d(n_inputs, n_outputs, 1) if n_inputs != n_outputs else None
+        )
+        self.relu = nn.ReLU()
+        self.init_weights()
+
+    def init_weights(self):
+        self.conv1.weight.data.normal_(0, 0.01)
+        self.conv2.weight.data.normal_(0, 0.01)
+        if self.downsample is not None:
+            self.downsample.weight.data.normal_(0, 0.01)
+
+    def forward(self, x):
+        out = self.net(x)
+        res = x if self.downsample is None else self.downsample(x)
+        return self.relu(out + res)
+
+
+class TCNLayer(nn.Module):
+    """Temporal Convolutional Networks layer.
+
+    Shaojie Bai et al. An Empirical Evaluation of Generic Convolutional and Recurrent Networks for Sequence Modeling.
+
+    This layer wraps the PyTorch TCN layer with masking and dropout support. It is
+    used in the TCN model. But it can also be used as a standalone layer.
+
+    Args:
+        input_dim: input feature size.
+        num_channels: int or list of ints. If int, the depth will be automatically decided by the max_seq_length. If list, number of channels in each layer.
+        max_seq_length: max sequence length. Used to compute the depth of the TCN.
+        kernel_size: kernel size of the TCN.
+        dropout: dropout rate. If non-zero, introduces a Dropout layer before each
+            TCN blocks. Default is 0.5.
+
+    Examples:
+        >>> from pyhealth.models import TCNLayer
+        >>> input = torch.randn(3, 128, 5)  # [batch size, sequence len, input_size]
+        >>> layer = TCNLayer(5, 64)
+        >>> outputs, last_outputs = layer(input)
+        >>> outputs.shape
+        torch.Size([3, 128, 64])
+        >>> last_outputs.shape
+        torch.Size([3, 64])
+    """
+
+    def __init__(
+        self,
+        input_dim: int,
+        num_channels: int = 128,
+        max_seq_length: int = 20,
+        kernel_size: int = 2,
+        dropout: float = 0.5,
+    ):
+        super(TCNLayer, self).__init__()
+        self.num_channels = num_channels
+
+        layers = []
+
+        # We compute automatically the depth based on the desired seq_length.
+        if isinstance(num_channels, int) and max_seq_length:
+            num_channels = [num_channels] * int(
+                np.ceil(np.log(max_seq_length / 2) / np.log(kernel_size))
+            )
+        elif isinstance(num_channels, int) and not max_seq_length:
+            raise Exception(
+                "a maximum sequence length needs to be provided if num_channels is int"
+            )
+        else:
+            pass
+
+        num_levels = len(num_channels)
+        for i in range(num_levels):
+            dilation_size = 2**i
+            in_channels = input_dim if i == 0 else num_channels[i - 1]
+            out_channels = num_channels[i]
+            layers += [
+                TemporalBlock(
+                    in_channels,
+                    out_channels,
+                    kernel_size,
+                    stride=1,
+                    dilation=dilation_size,
+                    padding=(kernel_size - 1) * dilation_size,
+                    dropout=dropout,
+                )
+            ]
+
+        self.network = nn.Sequential(*layers)
+
+    def forward(
+        self,
+        x: torch.tensor,
+        mask: Optional[torch.tensor] = None,
+    ) -> Tuple[torch.tensor, torch.tensor]:
+        """Forward propagation.
+
+        Args:
+            x: a tensor of shape [batch size, sequence len, input size].
+            mask: an optional tensor of shape [batch size, sequence len], where
+                1 indicates valid and 0 indicates invalid.
+
+        Returns:
+            last_out: a tensor of shape [batch size, hidden size], containing
+                the output features for the last time step.
+            out: a tensor of shape [batch size, sequence len, hidden size],
+                containing the output features for each time step.
+        """
+        out = self.network(x.permute(0, 2, 1)).permute(0, 2, 1)
+        last_out = get_last_visit(out, mask)
+        return last_out, out
+
+
+class TCN(BaseModel):
+    """Temporal Convolutional Networks model.
+
+    This model applies a separate TCN layer for each feature, and then concatenates
+    the final hidden states of each TCN layer. The concatenated hidden states are
+    then fed into a fully connected layer to make predictions.
+
+    Note:
+        We use separate TCN layers for different feature_keys.
+        Currently, we automatically support different input formats:
+            - code based input (need to use the embedding table later)
+            - float/int based value input
+        We follow the current convention for the TCN model:
+            - case 1. [code1, code2, code3, ...]
+                - we will assume the code follows the order; our model will encode
+                each code into a vector and apply TCN on the code level
+            - case 2. [[code1, code2]] or [[code1, code2], [code3, code4, code5], ...]
+                - we will assume the inner bracket follows the order; our model first
+                use the embedding table to encode each code into a vector and then use
+                average/mean pooling to get one vector for one inner bracket; then use
+                TCN one the braket level
+            - case 3. [[1.5, 2.0, 0.0]] or [[1.5, 2.0, 0.0], [8, 1.2, 4.5], ...]
+                - this case only makes sense when each inner bracket has the same length;
+                we assume each dimension has the same meaning; we run TCN directly
+                on the inner bracket level, similar to case 1 after embedding table
+            - case 4. [[[1.5, 2.0, 0.0]]] or [[[1.5, 2.0, 0.0], [8, 1.2, 4.5]], ...]
+                - this case only makes sense when each inner bracket has the same length;
+                we assume each dimension has the same meaning; we run TCN directly
+                on the inner bracket level, similar to case 2 after embedding table
+
+    Args:
+        dataset: the dataset to train the model. It is used to query certain
+            information such as the set of all tokens.
+        feature_keys:  list of keys in samples to use as features,
+            e.g. ["conditions", "procedures"].
+        label_key: key in samples to use as label (e.g., "drugs").
+        mode: one of "binary", "multiclass", or "multilabel".
+        embedding_dim: the embedding dimension. Default is 128.
+        num_channels: the number of channels in the TCN layer. Default is 128.
+        **kwargs: other parameters for the TCN layer.
+
+    Examples:
+        >>> from pyhealth.datasets import SampleEHRDataset
+        >>> samples = [
+        ...         {
+        ...             "patient_id": "patient-0",
+        ...             "visit_id": "visit-0",
+        ...             "list_codes": ["505800458", "50580045810", "50580045811"],  # NDC
+        ...             "list_vectors": [[1.0, 2.55, 3.4], [4.1, 5.5, 6.0]],
+        ...             "list_list_codes": [["A05B", "A05C", "A06A"], ["A11D", "A11E"]],  # ATC-4
+        ...             "list_list_vectors": [
+        ...                 [[1.8, 2.25, 3.41], [4.50, 5.9, 6.0]],
+        ...                 [[7.7, 8.5, 9.4]],
+        ...             ],
+        ...             "label": 1,
+        ...         },
+        ...         {
+        ...             "patient_id": "patient-0",
+        ...             "visit_id": "visit-1",
+        ...             "list_codes": [
+        ...                 "55154191800",
+        ...                 "551541928",
+        ...                 "55154192800",
+        ...                 "705182798",
+        ...                 "70518279800",
+        ...             ],
+        ...             "list_vectors": [[1.4, 3.2, 3.5], [4.1, 5.9, 1.7], [4.5, 5.9, 1.7]],
+        ...             "list_list_codes": [["A04A", "B035", "C129"]],
+        ...             "list_list_vectors": [
+        ...                 [[1.0, 2.8, 3.3], [4.9, 5.0, 6.6], [7.7, 8.4, 1.3], [7.7, 8.4, 1.3]],
+        ...             ],
+        ...             "label": 0,
+        ...         },
+        ...     ]
+        >>> dataset = SampleEHRDataset(samples=samples, dataset_name="test")
+        >>>
+        >>> from pyhealth.models import TCN
+        >>> model = TCN(
+        ...         dataset=dataset,
+        ...         feature_keys=[
+        ...             "list_codes",
+        ...             "list_vectors",
+        ...             "list_list_codes",
+        ...             "list_list_vectors",
+        ...         ],
+        ...         label_key="label",
+        ...         mode="binary",
+        ...     )
+        >>>
+        >>> from pyhealth.datasets import get_dataloader
+        >>> train_loader = get_dataloader(dataset, batch_size=2, shuffle=True)
+        >>> data_batch = next(iter(train_loader))
+        >>>
+        >>> ret = model(**data_batch)
+        >>> print(ret)
+        {'loss': tensor(0.7664, grad_fn=<BinaryCrossEntropyWithLogitsBackward0>), 'y_prob': tensor([[0.4714],
+                [0.4085]], grad_fn=<SigmoidBackward0>), 'y_true': tensor([[0.],
+                [1.]])}
+        >>>
+
+
+    """
+
+    def __init__(
+        self,
+        dataset: SampleEHRDataset,
+        feature_keys: List[str],
+        label_key: str,
+        mode: str,
+        embedding_dim: int = 128,
+        num_channels: int = 128,
+        **kwargs
+    ):
+        super(TCN, self).__init__(
+            dataset=dataset,
+            feature_keys=feature_keys,
+            label_key=label_key,
+            mode=mode,
+        )
+        self.embedding_dim = embedding_dim
+        self.num_channels = num_channels
+
+        if "input_dim" in kwargs:
+            raise ValueError("input_dim is determined by embedding_dim")
+
+        # the key of self.feat_tokenizers only contains the code based inputs
+        self.feat_tokenizers = {}
+        self.label_tokenizer = self.get_label_tokenizer()
+        # the key of self.embeddings only contains the code based inputs
+        self.embeddings = nn.ModuleDict()
+        # the key of self.linear_layers only contains the float/int based inputs
+        self.linear_layers = nn.ModuleDict()
+
+        for feature_key in self.feature_keys:
+            input_info = self.dataset.input_info[feature_key]
+            # sanity check
+            if input_info["type"] not in [str, float, int]:
+                raise ValueError(
+                    "TCN only supports str code, float and int as input types"
+                )
+            elif (input_info["type"] == str) and (input_info["dim"] not in [2, 3]):
+                raise ValueError(
+                    "TCN only supports 2-dim or 3-dim str code as input types"
+                )
+            elif (input_info["type"] in [float, int]) and (
+                input_info["dim"] not in [2, 3]
+            ):
+                raise ValueError(
+                    "TCN only supports 2-dim or 3-dim float and int as input types"
+                )
+            else:
+                pass
+            # for code based input, we need Type
+            # for float/int based input, we need Type, input_dim
+            self.add_feature_transform_layer(feature_key, input_info)
+
+        self.tcn = nn.ModuleDict()
+        for feature_key in feature_keys:
+            self.tcn[feature_key] = TCNLayer(
+                input_dim=embedding_dim, num_channels=num_channels, **kwargs
+            )
+        output_size = self.get_output_size(self.label_tokenizer)
+        self.fc = nn.Linear(len(self.feature_keys) * self.num_channels, output_size)
+
+    def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
+        """Forward propagation.
+
+        The label `kwargs[self.label_key]` is a list of labels for each patient.
+
+        Args:
+            **kwargs: keyword arguments for the model. The keys must contain
+                all the feature keys and the label key.
+
+        Returns:
+            A dictionary with the following keys:
+                loss: a scalar tensor representing the loss.
+                y_prob: a tensor representing the predicted probabilities.
+                y_true: a tensor representing the true labels.
+        """
+        patient_emb = []
+        for feature_key in self.feature_keys:
+            input_info = self.dataset.input_info[feature_key]
+            dim_, type_ = input_info["dim"], input_info["type"]
+
+            # for case 1: [code1, code2, code3, ...]
+            if (dim_ == 2) and (type_ == str):
+                x = self.feat_tokenizers[feature_key].batch_encode_2d(
+                    kwargs[feature_key]
+                )
+                # (patient, event)
+                x = torch.tensor(x, dtype=torch.long, device=self.device)
+                # (patient, event, embedding_dim)
+                x = self.embeddings[feature_key](x)
+                # (patient, event)
+                mask = torch.sum(x, dim=2) != 0
+
+            # for case 2: [[code1, code2], [code3, ...], ...]
+            elif (dim_ == 3) and (type_ == str):
+                x = self.feat_tokenizers[feature_key].batch_encode_3d(
+                    kwargs[feature_key]
+                )
+                # (patient, visit, event)
+                x = torch.tensor(x, dtype=torch.long, device=self.device)
+                # (patient, visit, event, embedding_dim)
+                x = self.embeddings[feature_key](x)
+                # (patient, visit, embedding_dim)
+                x = torch.sum(x, dim=2)
+                # (patient, visit)
+                mask = torch.sum(x, dim=2) != 0
+
+            # for case 3: [[1.5, 2.0, 0.0], ...]
+            elif (dim_ == 2) and (type_ in [float, int]):
+                x, mask = self.padding2d(kwargs[feature_key])
+                # (patient, event, values)
+                x = torch.tensor(x, dtype=torch.float, device=self.device)
+                # (patient, event, embedding_dim)
+                x = self.linear_layers[feature_key](x)
+                # (patient, event)
+                mask = mask.bool().to(self.device)
+
+            # for case 4: [[[1.5, 2.0, 0.0], [1.8, 2.4, 6.0]], ...]
+            elif (dim_ == 3) and (type_ in [float, int]):
+                x, mask = self.padding3d(kwargs[feature_key])
+                # (patient, visit, event, values)
+                x = torch.tensor(x, dtype=torch.float, device=self.device)
+                # (patient, visit, embedding_dim)
+                x = torch.sum(x, dim=2)
+                x = self.linear_layers[feature_key](x)
+                # (patient, event)
+                mask = mask[:, :, 0]
+                mask = mask.bool().to(self.device)
+
+            else:
+                raise NotImplementedError
+
+            x, _ = self.tcn[feature_key](x, mask)
+            patient_emb.append(x)
+
+        patient_emb = torch.cat(patient_emb, dim=1)
+        # (patient, label_size)
+        logits = self.fc(patient_emb)
+        # obtain y_true, loss, y_prob
+        y_true = self.prepare_labels(kwargs[self.label_key], self.label_tokenizer)
+        loss = self.get_loss_function()(logits, y_true)
+        y_prob = self.prepare_y_prob(logits)
+        return {
+            "loss": loss,
+            "y_prob": y_prob,
+            "y_true": y_true,
+        }
+
+
+if __name__ == "__main__":
+    from pyhealth.datasets import SampleEHRDataset
+
+    samples = [
+        {
+            "patient_id": "patient-0",
+            "visit_id": "visit-0",
+            # "single_vector": [1, 2, 3],
+            "list_codes": ["505800458", "50580045810", "50580045811"],  # NDC
+            "list_vectors": [[1.0, 2.55, 3.4], [4.1, 5.5, 6.0]],
+            "list_list_codes": [["A05B", "A05C", "A06A"], ["A11D", "A11E"]],  # ATC-4
+            "list_list_vectors": [
+                [[1.8, 2.25, 3.41], [4.50, 5.9, 6.0]],
+                [[7.7, 8.5, 9.4]],
+            ],
+            "label": 1,
+        },
+        {
+            "patient_id": "patient-0",
+            "visit_id": "visit-1",
+            # "single_vector": [1, 5, 8],
+            "list_codes": [
+                "55154191800",
+                "551541928",
+                "55154192800",
+                "705182798",
+                "70518279800",
+            ],
+            "list_vectors": [[1.4, 3.2, 3.5], [4.1, 5.9, 1.7], [4.5, 5.9, 1.7]],
+            "list_list_codes": [["A04A", "B035", "C129"]],
+            "list_list_vectors": [
+                [[1.0, 2.8, 3.3], [4.9, 5.0, 6.6], [7.7, 8.4, 1.3], [7.7, 8.4, 1.3]],
+            ],
+            "label": 0,
+        },
+    ]
+
+    # dataset
+    dataset = SampleEHRDataset(samples=samples, dataset_name="test")
+
+    # data loader
+    from pyhealth.datasets import get_dataloader
+
+    train_loader = get_dataloader(dataset, batch_size=2, shuffle=True)
+
+    # model
+    model = TCN(
+        dataset=dataset,
+        feature_keys=[
+            "list_codes",
+            "list_vectors",
+            "list_list_codes",
+            "list_list_vectors",
+        ],
+        label_key="label",
+        mode="binary",
+    )
+
+    # data batch
+    data_batch = next(iter(train_loader))
+
+    # try the model
+    ret = model(**data_batch)
+    print(ret)
+
+    # try loss backward
+    ret["loss"].backward()
diff --git a/pyhealth/models/utils.py b/pyhealth/models/utils.py
index 524966c4..d8e61bbb 100644
--- a/pyhealth/models/utils.py
+++ b/pyhealth/models/utils.py
@@ -29,11 +29,14 @@ def get_last_visit(hidden_states, mask):
     Returns:
         last_visit: [batch size, hidden_size]
     """
-    mask = mask.long()
-    last_visit = torch.sum(mask, 1) - 1
-    last_visit = last_visit.unsqueeze(-1)
-    last_visit = last_visit.expand(-1, hidden_states.shape[1] * hidden_states.shape[2])
-    last_visit = torch.reshape(last_visit, hidden_states.shape)
-    last_hidden_states = torch.gather(hidden_states, 1, last_visit)
-    last_hidden_state = last_hidden_states[:, 0, :]
-    return last_hidden_state
+    if mask is None:
+        return hidden_states[:, -1, :]
+    else:
+        mask = mask.long()
+        last_visit = torch.sum(mask, 1) - 1
+        last_visit = last_visit.unsqueeze(-1)
+        last_visit = last_visit.expand(-1, hidden_states.shape[1] * hidden_states.shape[2])
+        last_visit = torch.reshape(last_visit, hidden_states.shape)
+        last_hidden_states = torch.gather(hidden_states, 1, last_visit)
+        last_hidden_state = last_hidden_states[:, 0, :]
+        return last_hidden_state


Test patch:
diff --git a/pyhealth/models/test/test_gamenet.py b/pyhealth/models/test/test_gamenet.py
deleted file mode 100644
index f82dbb21..00000000
--- a/pyhealth/models/test/test_gamenet.py
+++ /dev/null
@@ -1,27 +0,0 @@
-from torch.utils.data import DataLoader
-
-from pyhealth.datasets import MIMIC3Dataset
-from pyhealth.models import GAMENet
-from pyhealth.tasks import drug_recommendation_mimic3_fn
-from pyhealth.datasets.utils import collate_fn_dict
-
-base_dataset = MIMIC3Dataset(
-    root="/srv/local/data/physionet.org/files/mimiciii/1.4",
-    tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
-    dev=True,
-    code_mapping={"NDC": "ATC"},
-    refresh_cache=False,
-)
-
-# visit level + multilabel
-sample_dataset = base_dataset.set_task(drug_recommendation_mimic3_fn)
-dataloader = DataLoader(
-    sample_dataset, batch_size=64, shuffle=True, collate_fn=collate_fn_dict
-)
-model = GAMENet(
-    dataset=sample_dataset,
-)
-model.to("cuda")
-batch = iter(dataloader).next()
-output = model(**batch)
-print(output["loss"])
diff --git a/pyhealth/models/test/test_general_model.py b/pyhealth/models/test/test_general_model.py
deleted file mode 100644
index e52cb19f..00000000
--- a/pyhealth/models/test/test_general_model.py
+++ /dev/null
@@ -1,157 +0,0 @@
-from torch.utils.data import DataLoader
-
-from pyhealth.datasets import MIMIC3Dataset
-from pyhealth.datasets.utils import collate_fn_dict
-
-# from pyhealth.models import CNN as Model
-# from pyhealth.models import RNN as Model
-from pyhealth.models import RETAIN as Model
-# from pyhealth.models import Transformer as Model
-
-
-def task_event(patient):
-    samples = []
-    for visit in patient:
-        conditions = visit.get_code_list(table="DIAGNOSES_ICD")
-        procedures = visit.get_code_list(table="PROCEDURES_ICD")
-        drugs = visit.get_code_list(table="PRESCRIPTIONS")
-        mortality_label = int(visit.discharge_status)
-        if len(conditions) * len(procedures) * len(drugs) == 0:
-            continue
-        samples.append(
-            {
-                "visit_id": visit.visit_id,
-                "patient_id": patient.patient_id,
-                "conditions": conditions,
-                "procedures": procedures,
-                "list_label": drugs,
-                "value_label": mortality_label,
-            }
-        )
-    return samples
-
-
-def task_visit(patient):
-    samples = []
-    for visit in patient:
-        conditions = visit.get_code_list(table="DIAGNOSES_ICD")
-        procedures = visit.get_code_list(table="PROCEDURES_ICD")
-        drugs = visit.get_code_list(table="PRESCRIPTIONS")
-        mortality_label = int(visit.discharge_status)
-        if len(conditions) * len(procedures) * len(drugs) == 0:
-            continue
-        samples.append(
-            {
-                "visit_id": visit.visit_id,
-                "patient_id": patient.patient_id,
-                "conditions": [conditions],
-                "procedures": [procedures],
-                "list_label": drugs,
-                "value_label": mortality_label,
-            }
-        )
-    return samples
-
-
-base_dataset = MIMIC3Dataset(
-    root="/srv/local/data/physionet.org/files/mimiciii/1.4",
-    tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
-    dev=True,
-    refresh_cache=False,
-)
-
-# event level + binary
-sample_dataset = base_dataset.set_task(task_event)
-dataloader = DataLoader(
-    sample_dataset, batch_size=64, shuffle=True, collate_fn=collate_fn_dict
-)
-model = Model(
-    dataset=sample_dataset,
-    feature_keys=["conditions", "procedures"],
-    label_key="value_label",
-    mode="binary",
-)
-model.to("cuda")
-batch = iter(dataloader).next()
-output = model(**batch)
-print(output["loss"])
-
-# visit level + binary
-sample_dataset = base_dataset.set_task(task_visit)
-dataloader = DataLoader(
-    sample_dataset, batch_size=64, shuffle=True, collate_fn=collate_fn_dict
-)
-model = Model(
-    dataset=sample_dataset,
-    feature_keys=["conditions", "procedures"],
-    label_key="value_label",
-    mode="binary",
-)
-model.to("cuda")
-batch = iter(dataloader).next()
-output = model(**batch)
-print(output["loss"])
-
-# event level + multiclass
-sample_dataset = base_dataset.set_task(task_event)
-dataloader = DataLoader(
-    sample_dataset, batch_size=64, shuffle=True, collate_fn=collate_fn_dict
-)
-model = Model(
-    dataset=sample_dataset,
-    feature_keys=["conditions", "procedures"],
-    label_key="value_label",
-    mode="multiclass",
-)
-model.to("cuda")
-batch = iter(dataloader).next()
-output = model(**batch)
-print(output["loss"])
-
-# visit level + multiclass
-sample_dataset = base_dataset.set_task(task_visit)
-dataloader = DataLoader(
-    sample_dataset, batch_size=64, shuffle=True, collate_fn=collate_fn_dict
-)
-model = Model(
-    dataset=sample_dataset,
-    feature_keys=["conditions", "procedures"],
-    label_key="value_label",
-    mode="multiclass",
-)
-model.to("cuda")
-batch = iter(dataloader).next()
-output = model(**batch)
-print(output["loss"])
-
-# event level + multilabel
-sample_dataset = base_dataset.set_task(task_event)
-dataloader = DataLoader(
-    sample_dataset, batch_size=64, shuffle=True, collate_fn=collate_fn_dict
-)
-model = Model(
-    dataset=sample_dataset,
-    feature_keys=["conditions", "procedures"],
-    label_key="list_label",
-    mode="multilabel",
-)
-model.to("cuda")
-batch = iter(dataloader).next()
-output = model(**batch)
-print(output["loss"])
-
-# visit level + multilabel
-sample_dataset = base_dataset.set_task(task_visit)
-dataloader = DataLoader(
-    sample_dataset, batch_size=64, shuffle=True, collate_fn=collate_fn_dict
-)
-model = Model(
-    dataset=sample_dataset,
-    feature_keys=["conditions", "procedures"],
-    label_key="list_label",
-    mode="multilabel",
-)
-model.to("cuda")
-batch = iter(dataloader).next()
-output = model(**batch)
-print(output["loss"])
diff --git a/pyhealth/models/test/test_micron.py b/pyhealth/models/test/test_micron.py
deleted file mode 100644
index 180e270d..00000000
--- a/pyhealth/models/test/test_micron.py
+++ /dev/null
@@ -1,27 +0,0 @@
-from torch.utils.data import DataLoader
-
-from pyhealth.datasets import MIMIC3Dataset
-from pyhealth.models import MICRON
-from pyhealth.tasks import drug_recommendation_mimic3_fn
-from pyhealth.datasets.utils import collate_fn_dict
-
-base_dataset = MIMIC3Dataset(
-    root="/srv/local/data/physionet.org/files/mimiciii/1.4",
-    tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
-    dev=True,
-    code_mapping={"NDC": "ATC"},
-    refresh_cache=False,
-)
-
-# visit level + multilabel
-sample_dataset = base_dataset.set_task(drug_recommendation_mimic3_fn)
-dataloader = DataLoader(
-    sample_dataset, batch_size=64, shuffle=True, collate_fn=collate_fn_dict
-)
-model = MICRON(
-    dataset=sample_dataset,
-)
-model.to("cuda")
-batch = iter(dataloader).next()
-output = model(**batch)
-print(output["loss"])
diff --git a/pyhealth/models/test/test_safedrug.py b/pyhealth/models/test/test_safedrug.py
deleted file mode 100644
index c732b2ac..00000000
--- a/pyhealth/models/test/test_safedrug.py
+++ /dev/null
@@ -1,27 +0,0 @@
-from torch.utils.data import DataLoader
-
-from pyhealth.datasets import MIMIC3Dataset
-from pyhealth.models import SafeDrug
-from pyhealth.tasks import drug_recommendation_mimic3_fn
-from pyhealth.datasets.utils import collate_fn_dict
-
-base_dataset = MIMIC3Dataset(
-    root="/srv/local/data/physionet.org/files/mimiciii/1.4",
-    tables=["DIAGNOSES_ICD", "PROCEDURES_ICD", "PRESCRIPTIONS"],
-    dev=True,
-    code_mapping={"NDC": "ATC"},
-    refresh_cache=False,
-)
-
-# visit level + multilabel
-sample_dataset = base_dataset.set_task(drug_recommendation_mimic3_fn)
-dataloader = DataLoader(
-    sample_dataset, batch_size=64, shuffle=True, collate_fn=collate_fn_dict
-)
-model = SafeDrug(
-    dataset=sample_dataset,
-)
-model.to("cuda")
-batch = iter(dataloader).next()
-output = model(**batch)
-print(output["loss"])