Repository: torcheeg/torcheeg (language: python)

Issue:
Adhere to standards



Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 5685971..0fc5c4e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -53,7 +53,7 @@ docs/source/generated
 docs/source/auto_examples/
 
 # Local experimental folders
-/tmp_in
-/tmp_out
-/dev
-/lightning_logs
\ No newline at end of file
+tmp_in/
+tmp_out/
+dev/
+lightning_logs/
\ No newline at end of file
diff --git a/README.md b/README.md
index 5025d68..96ad8b1 100644
--- a/README.md
+++ b/README.md
@@ -175,7 +175,7 @@ Here, `offline_transform` will only be called once when the dataset is initializ
 Next, we need to divide the dataset into a training set and a test set. In the field of EEG analysis, commonly used data partitioning methods include k-fold cross-validation and leave-one-out cross-validation. In this tutorial, we use k-fold cross-validation on the entire dataset (`KFold`) as an example of dataset splitting.
 
 ```python
-from torcheeg.datasets import KFoldGroupbyTrial
+from torcheeg.model_selection import KFoldGroupbyTrial
 
 k_fold = KFoldGroupbyTrial(n_splits=10,
                            split_path='./tmp_out/examples_pipeline/split',
@@ -261,6 +261,21 @@ If you would like to contribute new datasets, deep learning methods, and extensi
             <sub><b>Meekele_b</b></sub>
         </a>
     </td>
+    <td align="center">
+        <a href="https://github.com/dawin2015">
+            <img src="https://avatars.githubusercontent.com/u/12325204?v=4" width="100;" alt="dawin2015"/>
+            <br />
+            <sub><b>Liu-darong</b></sub>
+        </a>
+    </td>
+    <td align="center">
+        <a href="https://github.com/Nhix00">
+            <img src="https://avatars.githubusercontent.com/u/75837033?v=4" width="100;" alt="Nhix00"/>
+            <br />
+            <sub><b>Null</b></sub>
+        </a>
+    </td></tr>
+<tr>
     <td align="center">
         <a href="https://github.com/gitter-badger">
             <img src="https://avatars.githubusercontent.com/u/8518239?v=4" width="100;" alt="gitter-badger"/>
diff --git a/README.rst b/README.rst
index 9fe016d..3337626 100644
--- a/README.rst
+++ b/README.rst
@@ -299,7 +299,7 @@ this tutorial, we use k-fold cross-validation on the entire dataset
 
 .. code:: python
 
-   from torcheeg.datasets import KFoldGroupbyTrial
+   from torcheeg.model_selection import KFoldGroupbyTrial
 
    k_fold = KFoldGroupbyTrial(n_splits=10,
                               split_path='./tmp_out/examples_pipeline/split',
diff --git a/docs/requirements-windows.txt b/docs/requirements-windows.txt
new file mode 100644
index 0000000..2da3d4a
--- /dev/null
+++ b/docs/requirements-windows.txt
@@ -0,0 +1,28 @@
+https://download.pytorch.org/whl/cpu/torch-1.13.0%2Bcpu-cp38-cp38-win_amd64.whl
+https://data.pyg.org/whl/torch-1.13.0%2Bcpu/torch_scatter-2.0.9-cp38-cp38-win_amd64.whl
+https://data.pyg.org/whl/torch-1.13.0%2Bcpu/torch_sparse-0.6.17%2Bpt113cpu-cp38-cp38-win_amd64.whl
+mne
+mne_connectivity
+sphinx==6.2.1
+sphinx_gallery
+docutils
+pylons-sphinx-themes
+pylons_sphinx_latesturl
+repoze.sphinx.autointerface
+sphinxcontrib-autoprogram
+torchmetrics
+tqdm
+numpy
+pandas
+scipy
+scikit-learn
+lmdb
+einops
+torch_geometric
+xmltodict
+networkx
+PyWavelets
+spectrum
+moabb
+pytorch-lightning
+-e git+https://github.com/pytorch/pytorch_sphinx_theme.git#egg=pytorch_sphinx_theme
\ No newline at end of file
diff --git a/docs/source/torcheeg.datasets.rst b/docs/source/torcheeg.datasets.rst
index 6770487..1de1431 100644
--- a/docs/source/torcheeg.datasets.rst
+++ b/docs/source/torcheeg.datasets.rst
@@ -37,6 +37,16 @@ Personal Identification Datasets
 
     M3CVDataset
 
+Motor Imagery Datasets
+----------------------------------
+
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+    :template: datasettemplate.rst
+
+    BCICIV2aDataset
+
 Mother of all BCI Benchmarks
 ---------------------------------------------------------
 
diff --git a/docs/source/torcheeg.models.rst b/docs/source/torcheeg.models.rst
index b13e75c..b8a0dd3 100644
--- a/docs/source/torcheeg.models.rst
+++ b/docs/source/torcheeg.models.rst
@@ -19,6 +19,7 @@ Convolutional Neural Networks
    EEGNet
    FBCCNN
    FBCNet
+   FBMSNet
    MTCNN
    STNet
    TSCeption
@@ -61,6 +62,7 @@ Transformer
    ArjunViT
    VanillaTransformer
    ViT
+   ATCNet
 
 Generative Adversarial Network
 ----------------------------------
@@ -73,6 +75,8 @@ Generative Adversarial Network
    BDiscriminator
    BCGenerator
    BCDiscriminator
+   EEGfuseNet
+   EFDiscriminator   
 
 Variational Auto Encoder
 ----------------------------------
diff --git a/docs/source/torcheeg.trainers.rst b/docs/source/torcheeg.trainers.rst
index d40b2bd..f88d9b2 100644
--- a/docs/source/torcheeg.trainers.rst
+++ b/docs/source/torcheeg.trainers.rst
@@ -34,6 +34,23 @@ The individual differences and nonstationary of EEG signals make it difficult fo
    JANTrainer
    ADATrainer
    DANNTrainer
+   CenterLossTrainer
+
+Imbalance Learning for Classification
+----------------------------------
+
+EEG emotion datasets have the problem of sample class imbalance, and imbalance learning can be used to solve the class imbalance problem in emotion recognition tasks.
+
+.. autosummary::
+   :toctree: imbalance
+   :nosignatures:
+   :template: trainertemplate.rst
+
+   LALossTrainer
+   LDAMLossTrainer
+   EQLossTrainer
+   FocalLossTrainer
+   WCELossTrainer
 
 EEG Generation
 ----------------------------------------
diff --git a/docs/source/torcheeg.transforms.rst b/docs/source/torcheeg.transforms.rst
index 36559ea..7247001 100644
--- a/docs/source/torcheeg.transforms.rst
+++ b/docs/source/torcheeg.transforms.rst
@@ -30,6 +30,7 @@ Numpy-based Transforms
     :template: transformtemplate.rst
 
     CWTSpectrum
+    BandSignal
     BandDifferentialEntropy
     BandPowerSpectralDensity
     BandMeanAbsoluteDeviation
diff --git a/examples/examples_deap_cwgangp.py b/examples/examples_deap_cwgangp.py
index 4720fe3..6c06e92 100644
--- a/examples/examples_deap_cwgangp.py
+++ b/examples/examples_deap_cwgangp.py
@@ -57,7 +57,7 @@
 # testing.
 # 
 
-from torcheeg.datasets import KFoldGroupbyTrial
+from torcheeg.model_selection import KFoldGroupbyTrial
 
 k_fold = KFoldGroupbyTrial(n_splits=5,
                            split_path='./examples_deap_cwgangp/split',
diff --git a/examples/examples_models.py b/examples/examples_models.py
index a18a466..10c682f 100644
--- a/examples/examples_models.py
+++ b/examples/examples_models.py
@@ -118,6 +118,37 @@
 pred = model(eeg)
 
 
+######################################################################
+# Some studies have shown that attention based models have achieved good 
+# classification performance in EEG, such as Altaheri et al.'s ATCNet, 
+# which uses moving windows in the model structure and utilizes 
+# multiheadattention to process data within the window.This model achieved 
+# excellent results in data set 2a of the BCI Competition IV.
+#
+
+from torcheeg.models import ATCNet
+from torcheeg.datasets import BCICIV2aDataset
+from torcheeg import transforms
+
+dataset = BCICIV2aDataset(io_path=f'./bciciv_2a',
+                              root_path='./BCICIV_2a_mat',
+                              online_transform=transforms.Compose([
+                                  transforms.To2d(),
+                                  transforms.ToTensor()
+                              ]),
+                              label_transform=transforms.Compose([
+                                  transforms.Select('label'),
+                                  transforms.Lambda(lambda x: x - 1)
+                              ]))                             
+model = ATCNet(num_classes=4,
+               num_windows=3,
+               num_electrodes=22,
+               chunk_size=1750)
+x = dataset[0][0]
+x = torch.unsqueeze(x,dim=0)
+pred = model(x)
+
+
 ######################################################################
 # Generative Models
 # ----------------------------------------------
@@ -203,3 +234,33 @@
 t = torch.randint(low=1, high=1000, size=(2, ))
 y = torch.randint(low=0, high=2, size=(1, ))
 fake_X = unet(mock_eeg, t, y)
+
+
+######################################################################
+# Eegtorch provides other types of models, such as eegfusenet, which combines 
+# the functions of EEG encoding and generating new samples. At the same time, 
+# eefusenet is an unsupervised learning model that can extract deep feature 
+# encoding from input EEG signals and ultimately generate similar new samples. 
+# Eegfusenet uses a approach similar to traditional gan models to identify 
+# whether samples is real: EFDiscriminator, which ultimately improves the quality 
+# of sample generation through eegfuset after adversarial training.
+#
+
+from torcheeg.models import EEGfuseNet,EFDiscriminator
+
+fusenet = EEGfuseNet(in_channels=1,
+                     num_electrodes=32,
+                     hid_channels_gru=16,
+                     num_layers_gru= 1,
+                     hid_channels_cnn=1,
+                     chunk_size=384)
+eeg = torch.randn(2,1, 32, 384) 
+# simply input the EEG signal to output generated samples and deep fusion codes
+fake_X,deep_code = fusenet(eeg)
+
+discriminator = EFDiscriminator(in_channels=1,
+                                num_electrodes=32,
+                                hid_channels_cnn=1,
+                                chunk_size=384)
+p_real = discriminator(eeg)
+p_fake = discriminator(fake_X)
\ No newline at end of file
diff --git a/examples/examples_pipeline.py b/examples/examples_pipeline.py
index 08b52c3..b07f84b 100644
--- a/examples/examples_pipeline.py
+++ b/examples/examples_pipeline.py
@@ -61,7 +61,7 @@
 # testing.
 # 
 
-from torcheeg.datasets import KFoldGroupbyTrial
+from torcheeg.model_selection import KFoldGroupbyTrial
 
 k_fold = KFoldGroupbyTrial(n_splits=10,
                            split_path='./examples_pipeline/split',
diff --git a/torcheeg/datasets/constants/motor_imagery/bciciv_2a.py b/torcheeg/datasets/constants/motor_imagery/bciciv_2a.py
new file mode 100644
index 0000000..fef34ef
--- /dev/null
+++ b/torcheeg/datasets/constants/motor_imagery/bciciv_2a.py
@@ -0,0 +1,3 @@
+BCICIV2A_CHANNEL_LIST = [
+    'FZ', '0', '1', '2', '3', '4', '5', 'C3', '6', 'CZ', '7', 'C4', '8', '9', '10', '11', '12', '13', '14', 'PZ', '15', '16'
+]
\ No newline at end of file
diff --git a/torcheeg/datasets/constants/motor_imagery/physionet_mi.py b/torcheeg/datasets/constants/motor_imagery/physionet_mi.py
deleted file mode 100644
index ec876f2..0000000
--- a/torcheeg/datasets/constants/motor_imagery/physionet_mi.py
+++ /dev/null
@@ -1,110 +0,0 @@
-from ..region_1020 import (FRONTAL_REGION_LIST, GENERAL_REGION_LIST,
-                           HEMISPHERE_REGION_LIST)
-from ..standard_1020 import STANDARD_1020_CHANNEL_LOCATION_DICT
-from ..utils import (format_adj_matrix_from_adj_list,
-                     format_adj_matrix_from_standard,
-                     format_channel_location_dict, format_region_channel_list)
-
-PHYSIONETMI_CHANNEL_LIST = [
-    'FC5', 'FC3', 'FC1', 'FCZ', 'FC2', 'FC4', 'FC6', 'C5', 'C3', 'C1', 'CZ',
-    'C2', 'C4', 'C6', 'CP5', 'CP3', 'CP1', 'CPZ', 'CP2', 'CP4', 'CP6', 'FP1',
-    'FPZ', 'FP2', 'AF7', 'AF3', 'AFZ', 'AF4', 'AF8', 'F7', 'F5', 'F3', 'F1',
-    'FZ', 'F2', 'F4', 'F6', 'F8', 'FT7', 'FT8', 'T7', 'T8', 'T9', 'T10', 'TP7',
-    'TP8', 'P7', 'P5', 'P3', 'P1', 'PZ', 'P2', 'P4', 'P6', 'P8', 'PO7', 'PO3',
-    'POZ', 'PO4', 'PO8', 'O1', 'OZ', 'O2', 'IZ'
-]
-
-PHYSIONETMI_LOCATION_LIST = [
-    ['-', '-', '-', 'FP1', 'FPZ', 'FP2', '-', '-', '-'],
-    ['-', 'AF7', '-', 'AF3', 'AFZ', 'AF4', '-', 'AF8', '-'],
-    ['F7', 'F5', 'F3', 'F1', 'FZ', 'F2', 'F4', 'F6', 'F8'],
-    ['FT7', 'FC5', 'FC3', 'FC1', 'FCZ', 'FC2', 'FC4', 'FC6', 'FT8'],
-    ['T7', 'C5', 'C3', 'C1', 'CZ', 'C2', 'C4', 'C6', 'T8'],
-    ['TP7', 'CP5', 'CP3', 'CP1', 'CPZ', 'CP2', 'CP4', 'CP6', 'TP8'],
-    ['P7', 'P5', 'P3', 'P1', 'PZ', 'P2', 'P4', 'P6', 'P8'],
-    ['-', 'PO7', '-', 'PO3', 'POZ', 'PO4', '-', 'PO8', '-'],
-    ['-', '-', '-', 'O1', 'OZ', 'O2', '-', '-', '-']
-]
-# without {'T10', 'T9', 'IZ'}
-
-PHYSIONETMI_CHANNEL_LOCATION_DICT = format_channel_location_dict(
-    PHYSIONETMI_CHANNEL_LIST, PHYSIONETMI_LOCATION_LIST)
-
-PHYSIONETMI_ADJACENCY_LIST = {
-    'FP1': ['AF3', 'FPZ', 'AFZ'],
-    'FPZ': ['AFZ', 'FP1', 'FP2', 'AF4', 'AF3'],
-    'FP2': ['AF4', 'FPZ', 'AFZ'],
-    'AF7': ['F5', 'F3', 'F7'],
-    'AF3': ['FP1', 'F1', 'AFZ', 'FPZ', 'FZ', 'F3'],
-    'AFZ': ['FPZ', 'FZ', 'AF3', 'AF4', 'FP2', 'F2', 'FP1', 'F1'],
-    'AF4': ['FP2', 'F2', 'AFZ', 'F4', 'FPZ', 'FZ'],
-    'AF8': ['F6', 'F8', 'F4'],
-    'F7': ['FT7', 'F5', 'AF7', 'FC5'],
-    'F5': ['AF7', 'FC5', 'F7', 'F3', 'FC3', 'FT7'],
-    'F3': ['FC3', 'F5', 'F1', 'AF3', 'FC1', 'AF7', 'FC5'],
-    'F1': ['AF3', 'FC1', 'F3', 'FZ', 'AFZ', 'FCZ', 'FC3'],
-    'FZ': ['AFZ', 'FCZ', 'F1', 'F2', 'AF4', 'FC2', 'AF3', 'FC1'],
-    'F2': ['AF4', 'FC2', 'FZ', 'F4', 'FC4', 'AFZ', 'FCZ'],
-    'F4': ['FC4', 'F2', 'F6', 'AF8', 'FC6', 'AF4', 'FC2'],
-    'F6': ['AF8', 'FC6', 'F4', 'F8', 'FT8', 'FC4'],
-    'F8': ['FT8', 'F6', 'AF8', 'FC6'],
-    'FT7': ['F7', 'T7', 'FC5', 'F5', 'C5'],
-    'FC5': ['F5', 'C5', 'FT7', 'FC3', 'F3', 'C3', 'F7', 'T7'],
-    'FC3': ['F3', 'C3', 'FC5', 'FC1', 'F1', 'C1', 'F5', 'C5'],
-    'FC1': ['F1', 'C1', 'FC3', 'FCZ', 'FZ', 'CZ', 'F3', 'C3'],
-    'FCZ': ['FZ', 'CZ', 'FC1', 'FC2', 'F2', 'C2', 'F1', 'C1'],
-    'FC2': ['F2', 'C2', 'FCZ', 'FC4', 'F4', 'C4', 'FZ', 'CZ'],
-    'FC4': ['F4', 'C4', 'FC2', 'FC6', 'F6', 'C6', 'F2', 'C2'],
-    'FC6': ['F6', 'C6', 'FC4', 'FT8', 'F8', 'T8', 'F4', 'C4'],
-    'FT8': ['F8', 'T8', 'FC6', 'F6', 'C6'],
-    'T7': ['FT7', 'TP7', 'C5', 'FC5', 'CP5'],
-    'C5': ['FC5', 'CP5', 'T7', 'C3', 'FC3', 'CP3', 'FT7', 'TP7'],
-    'C3': ['FC3', 'CP3', 'C5', 'C1', 'FC1', 'CP1', 'FC5', 'CP5'],
-    'C1': ['FC1', 'CP1', 'C3', 'CZ', 'FCZ', 'CPZ', 'FC3', 'CP3'],
-    'CZ': ['FCZ', 'CPZ', 'C1', 'C2', 'FC2', 'CP2', 'FC1', 'CP1'],
-    'C2': ['FC2', 'CP2', 'CZ', 'C4', 'FC4', 'CP4', 'FCZ', 'CPZ'],
-    'C4': ['FC4', 'CP4', 'C2', 'C6', 'FC6', 'CP6', 'FC2', 'CP2'],
-    'C6': ['FC6', 'CP6', 'C4', 'T8', 'FT8', 'TP8', 'FC4', 'CP4'],
-    'T8': ['FT8', 'TP8', 'C6', 'FC6', 'CP6'],
-    'TP7': ['T7', 'P7', 'CP5', 'C5', 'P5'],
-    'CP5': ['C5', 'P5', 'TP7', 'CP3', 'C3', 'P3', 'T7', 'P7'],
-    'CP3': ['C3', 'P3', 'CP5', 'CP1', 'C1', 'P1', 'C5', 'P5'],
-    'CP1': ['C1', 'P1', 'CP3', 'CPZ', 'CZ', 'PZ', 'C3', 'P3'],
-    'CPZ': ['CZ', 'PZ', 'CP1', 'CP2', 'C2', 'P2', 'C1', 'P1'],
-    'CP2': ['C2', 'P2', 'CPZ', 'CP4', 'C4', 'P4', 'CZ', 'PZ'],
-    'CP4': ['C4', 'P4', 'CP2', 'CP6', 'C6', 'P6', 'C2', 'P2'],
-    'CP6': ['C6', 'P6', 'CP4', 'TP8', 'T8', 'P8', 'C4', 'P4'],
-    'TP8': ['T8', 'P8', 'CP6', 'C6', 'P6'],
-    'P7': ['TP7', 'P5', 'CP5', 'PO7'],
-    'P5': ['CP5', 'PO7', 'P7', 'P3', 'CP3', 'TP7'],
-    'P3': ['CP3', 'P5', 'P1', 'CP1', 'PO3', 'CP5', 'PO7'],
-    'P1': ['CP1', 'PO3', 'P3', 'PZ', 'CPZ', 'POZ', 'CP3'],
-    'PZ': ['CPZ', 'POZ', 'P1', 'P2', 'CP2', 'PO4', 'CP1', 'PO3'],
-    'P2': ['CP2', 'PO4', 'PZ', 'P4', 'CP4', 'CPZ', 'POZ'],
-    'P4': ['CP4', 'P2', 'P6', 'CP6', 'PO8', 'CP2', 'PO4'],
-    'P6': ['CP6', 'PO8', 'P4', 'P8', 'TP8', 'CP4'],
-    'P8': ['TP8', 'P6', 'CP6', 'PO8'],
-    'PO7': ['P5', 'P3', 'P7'],
-    'PO3': ['P1', 'O1', 'POZ', 'PZ', 'OZ', 'P3'],
-    'POZ': ['PZ', 'OZ', 'PO3', 'PO4', 'P2', 'O2', 'P1', 'O1'],
-    'PO4': ['P2', 'O2', 'POZ', 'P4', 'PZ', 'OZ'],
-    'PO8': ['P6', 'P8', 'P4'],
-    'O1': ['PO3', 'OZ', 'POZ'],
-    'OZ': ['POZ', 'O1', 'O2', 'PO4', 'PO3'],
-    'O2': ['PO4', 'OZ', 'POZ']
-}
-
-PHYSIONETMI_ADJACENCY_MATRIX = format_adj_matrix_from_adj_list(
-    PHYSIONETMI_CHANNEL_LIST, PHYSIONETMI_ADJACENCY_LIST)
-
-PHYSIONETMI_STANDARD_ADJACENCY_MATRIX = format_adj_matrix_from_standard(
-    PHYSIONETMI_CHANNEL_LIST,
-    STANDARD_1020_CHANNEL_LOCATION_DICT,
-    delta=0.00035)
-
-PHYSIONETMI_GENERAL_REGION_LIST = format_region_channel_list(
-    PHYSIONETMI_CHANNEL_LIST, GENERAL_REGION_LIST)
-PHYSIONETMI_FRONTAL_REGION_LIST = format_region_channel_list(
-    PHYSIONETMI_CHANNEL_LIST, FRONTAL_REGION_LIST)
-PHYSIONETMI_HEMISPHERE_REGION_LIST = format_region_channel_list(
-    PHYSIONETMI_CHANNEL_LIST, HEMISPHERE_REGION_LIST)
diff --git a/torcheeg/datasets/moabb/moabb_dataset.py b/torcheeg/datasets/moabb/moabb_dataset.py
index e39321e..62bdc07 100644
--- a/torcheeg/datasets/moabb/moabb_dataset.py
+++ b/torcheeg/datasets/moabb/moabb_dataset.py
@@ -171,25 +171,12 @@ def process_record(file: Any = None,
                     start_at = start_at + step
                     end_at = start_at + chunk_size
 
-                    if after_trial is not None:
-                        trial_queue.append({
-                            'eeg': t_eeg,
-                            'key': clip_id,
-                            'info': record_info
-                        })
-                    else:
-                        yield {
+                    yield {
                             'eeg': t_eeg,
                             'key': clip_id,
                             'info': record_info
                         }
 
-            if len(trial_queue) and after_trial is not None:
-                trial_queue = after_trial(trial_queue)
-                for obj in trial_queue:
-                    assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
-                    yield obj
-
     def set_records(self,
                     dataset: _MOABBDataset,
                     io_path: str = './io/moabb',
diff --git a/torcheeg/datasets/module/__init__.py b/torcheeg/datasets/module/__init__.py
index f77f3e0..322dc2b 100644
--- a/torcheeg/datasets/module/__init__.py
+++ b/torcheeg/datasets/module/__init__.py
@@ -1,6 +1,8 @@
 from .base_dataset import BaseDataset
-from .emotion_recognition import *
 from .folder_dataset import FolderDataset
 from .csv_folder_dataset import CSVFolderDataset
-from .personal_identification import *
-from .sspev import *
\ No newline at end of file
+
+from .sspev import *
+from .motor_imagery import *
+from .emotion_recognition import *
+from .personal_identification import *
\ No newline at end of file
diff --git a/torcheeg/datasets/module/base_dataset.py b/torcheeg/datasets/module/base_dataset.py
index a428057..426a569 100644
--- a/torcheeg/datasets/module/base_dataset.py
+++ b/torcheeg/datasets/module/base_dataset.py
@@ -17,6 +17,7 @@
 
 
 class MockLock():
+
     def __enter__(self):
         pass
 
@@ -25,6 +26,7 @@ def __exit__(self, exc_type, exc_value, traceback):
 
 
 class BaseDataset(Dataset):
+
     def __init__(self,
                  io_path: str = None,
                  io_size: int = 10485760,
@@ -97,15 +99,17 @@ def __init__(self,
             self.eeg_io_router, self.info = self.get_pointer(io_path=io_path,
                                                              io_size=io_size,
                                                              io_mode=io_mode)
-            # catch the exception
-            try:
-                self.post_process_record(after_trial=after_trial,
-                                         after_session=after_session,
-                                         after_subject=after_subject)
-            except Exception as e:
-                # shutil to delete the database
-                shutil.rmtree(self.io_path)
-                raise e
+
+            if self.after_trial is not None or self.after_session is not None or self.after_subject is not None:
+                # catch the exception
+                try:
+                    self.post_process_record(after_trial=after_trial,
+                                             after_session=after_session,
+                                             after_subject=after_subject)
+                except Exception as e:
+                    # shutil to delete the database
+                    shutil.rmtree(self.io_path)
+                    raise e
         else:
             print(
                 f'dataset already exists at path {self.io_path}, reading from path...'
diff --git a/torcheeg/datasets/module/emotion_recognition/amigos.py b/torcheeg/datasets/module/emotion_recognition/amigos.py
index 3d89aca..74ab92c 100644
--- a/torcheeg/datasets/module/emotion_recognition/amigos.py
+++ b/torcheeg/datasets/module/emotion_recognition/amigos.py
@@ -1,16 +1,10 @@
 import os
-import logging
 import re
 from typing import Any, Callable, Dict, List, Tuple, Union
 
 import scipy.io as scio
-
-from ...constants.emotion_recognition.amigos import (
-    AMIGOS_ADJACENCY_MATRIX, AMIGOS_CHANNEL_LOCATION_DICT)
 from ..base_dataset import BaseDataset
 
-log = logging.getLogger(__name__)
-
 class AMIGOSDataset(BaseDataset):
     r'''
     A dataset for Multimodal research of affect, personality traits and mood on Individuals and GrOupS (AMIGOS). This class generates training samples and test samples according to the given parameters, and caches the generated results in a unified input and output format (IO). The relevant information of the dataset is as follows:
@@ -129,8 +123,6 @@ class AMIGOSDataset(BaseDataset):
         verbose (bool): Whether to display logs during processing, such as progress bars, etc. (default: :obj:`True`)
         in_memory (bool): Whether to load the entire dataset into memory. If :obj:`in_memory` is set to True, then the first time an EEG sample is read, the entire dataset is loaded into memory for subsequent retrieval. Otherwise, the dataset is stored on disk to avoid the out-of-memory problem. (default: :obj:`False`)
     '''
-    channel_location_dict = AMIGOS_CHANNEL_LOCATION_DICT
-    adjacency_matrix = AMIGOS_ADJACENCY_MATRIX
 
     def __init__(self,
                  root_path: str = './data_preprocessed',
@@ -194,7 +186,6 @@ def process_record(file: Any = None,
                    baseline_chunk_size: int = 128,
                    before_trial: Union[None, Callable] = None,
                    offline_transform: Union[None, Callable] = None,
-                   after_trial: Union[None, Callable] = None,
                    **kwargs):
         file_name = file  # an element from file name list
 
diff --git a/torcheeg/datasets/module/emotion_recognition/bci2022.py b/torcheeg/datasets/module/emotion_recognition/bci2022.py
index 6122576..9c83656 100644
--- a/torcheeg/datasets/module/emotion_recognition/bci2022.py
+++ b/torcheeg/datasets/module/emotion_recognition/bci2022.py
@@ -2,11 +2,6 @@
 from typing import Any, Callable, Dict, Tuple, Union
 
 import joblib
-
-from torcheeg.io import EEGSignalIO, MetaInfoIO
-
-from ...constants.emotion_recognition.bci2022 import (
-    BCI2022_ADJACENCY_MATRIX, BCI2022_CHANNEL_LOCATION_DICT)
 from ..base_dataset import BaseDataset
 
 FIRST_BATCH_CHANNEL = [
@@ -198,8 +193,6 @@ class BCI2022Dataset(BaseDataset):
         verbose (bool): Whether to display logs during processing, such as progress bars, etc. (default: :obj:`True`)
         in_memory (bool): Whether to load the entire dataset into memory. If :obj:`in_memory` is set to True, then the first time an EEG sample is read, the entire dataset is loaded into memory for subsequent retrieval. Otherwise, the dataset is stored on disk to avoid the out-of-memory problem. (default: :obj:`False`)    
     '''
-    channel_location_dict = BCI2022_CHANNEL_LOCATION_DICT
-    adjacency_matrix = BCI2022_ADJACENCY_MATRIX
 
     def __init__(self,
                  root_path: str = './2022EmotionPublic/TrainSet/',
@@ -246,7 +239,6 @@ def process_record(file: Any = None,
                    channel_num: int = 30,
                    before_trial: Union[None, Callable] = None,
                    offline_transform: Union[None, Callable] = None,
-                   after_trial: Union[None, Callable] = None,
                    **kwargs):
 
         file_name = os.path.basename(file)  # an element from file name list
diff --git a/torcheeg/datasets/module/emotion_recognition/deap.py b/torcheeg/datasets/module/emotion_recognition/deap.py
index 222a0d2..d62ec90 100644
--- a/torcheeg/datasets/module/emotion_recognition/deap.py
+++ b/torcheeg/datasets/module/emotion_recognition/deap.py
@@ -1,7 +1,7 @@
 import os
 import pickle as pkl
-from typing import Any, Callable, Dict, Tuple, Union
 
+from typing import Any, Callable, Dict, Tuple, Union
 from ..base_dataset import BaseDataset
 
 
@@ -177,7 +177,6 @@ def process_record(file: Any = None,
                    num_baseline: int = 3,
                    baseline_chunk_size: int = 128,
                    before_trial: Union[None, Callable] = None,
-                   after_trial: Union[Callable, None] = None,
                    offline_transform: Union[None, Callable] = None,
                    **kwargs):
         file_name = file  # an element from file name list
diff --git a/torcheeg/datasets/module/emotion_recognition/dreamer.py b/torcheeg/datasets/module/emotion_recognition/dreamer.py
index 9334075..0642d91 100644
--- a/torcheeg/datasets/module/emotion_recognition/dreamer.py
+++ b/torcheeg/datasets/module/emotion_recognition/dreamer.py
@@ -1,10 +1,8 @@
 import os
 from typing import Callable, Dict, Tuple, Union, Any
+
 import scipy.io as scio
-from ...constants.emotion_recognition.dreamer import (
-    DREAMER_ADJACENCY_MATRIX, DREAMER_CHANNEL_LOCATION_DICT)
 from ..base_dataset import BaseDataset
-from torcheeg.io import EEGSignalIO, MetaInfoIO
 
 
 class DREAMERDataset(BaseDataset):
@@ -118,8 +116,6 @@ class DREAMERDataset(BaseDataset):
         verbose (bool): Whether to display logs during processing, such as progress bars, etc. (default: :obj:`True`)
         in_memory (bool): Whether to load the entire dataset into memory. If :obj:`in_memory` is set to True, then the first time an EEG sample is read, the entire dataset is loaded into memory for subsequent retrieval. Otherwise, the dataset is stored on disk to avoid the out-of-memory problem. (default: :obj:`False`)       
     '''
-    channel_location_dict = DREAMER_CHANNEL_LOCATION_DICT
-    adjacency_matrix = DREAMER_ADJACENCY_MATRIX
 
     def __init__(self,
                  mat_path: str = './DREAMER.mat',
@@ -177,7 +173,6 @@ def process_record(file: Any = None,
                    baseline_chunk_size: int = 128,
                    before_trial: Union[None, Callable] = None,
                    offline_transform: Union[None, Callable] = None,
-                   after_trial: Union[None, Callable] = None,
                    **kwargs):
         subject = file
         mat_data = scio.loadmat(mat_path,
diff --git a/torcheeg/datasets/module/emotion_recognition/mahnob.py b/torcheeg/datasets/module/emotion_recognition/mahnob.py
index 7ff8135..6e172cc 100644
--- a/torcheeg/datasets/module/emotion_recognition/mahnob.py
+++ b/torcheeg/datasets/module/emotion_recognition/mahnob.py
@@ -2,13 +2,9 @@
 import json
 import os
 from typing import Any, Callable, Dict, Tuple, Union
-from torcheeg.io import EEGSignalIO, MetaInfoIO
 
 import mne
 import xmltodict
-
-from ...constants.emotion_recognition.mahnob import (
-    MAHNOB_ADJACENCY_MATRIX, MAHNOB_CHANNEL_LOCATION_DICT)
 from ..base_dataset import BaseDataset
 
 
@@ -136,8 +132,6 @@ class MAHNOBDataset(BaseDataset):
         verbose (bool): Whether to display logs during processing, such as progress bars, etc. (default: :obj:`True`)
         in_memory (bool): Whether to load the entire dataset into memory. If :obj:`in_memory` is set to True, then the first time an EEG sample is read, the entire dataset is loaded into memory for subsequent retrieval. Otherwise, the dataset is stored on disk to avoid the out-of-memory problem. (default: :obj:`False`)    
     '''
-    channel_location_dict = MAHNOB_CHANNEL_LOCATION_DICT
-    adjacency_matrix = MAHNOB_ADJACENCY_MATRIX
 
     def __init__(self,
                  root_path: str = './Sessions',
diff --git a/torcheeg/datasets/module/emotion_recognition/mped_feature.py b/torcheeg/datasets/module/emotion_recognition/mped_feature.py
index 424a2e1..4dd4894 100644
--- a/torcheeg/datasets/module/emotion_recognition/mped_feature.py
+++ b/torcheeg/datasets/module/emotion_recognition/mped_feature.py
@@ -3,11 +3,6 @@
 
 import numpy as np
 import scipy.io as scio
-
-from torcheeg.io import EEGSignalIO, MetaInfoIO
-
-from ...constants.emotion_recognition.mped import (MPED_ADJACENCY_MATRIX,
-                                                   MPED_CHANNEL_LOCATION_DICT)
 from ..base_dataset import BaseDataset
 
 
@@ -100,8 +95,6 @@ class MPEDFeatureDataset(BaseDataset):
         verbose (bool): Whether to display logs during processing, such as progress bars, etc. (default: :obj:`True`)
         in_memory (bool): Whether to load the entire dataset into memory. If :obj:`in_memory` is set to True, then the first time an EEG sample is read, the entire dataset is loaded into memory for subsequent retrieval. Otherwise, the dataset is stored on disk to avoid the out-of-memory problem. (default: :obj:`False`)    
     '''
-    channel_location_dict = MPED_CHANNEL_LOCATION_DICT
-    adjacency_matrix = MPED_ADJACENCY_MATRIX
 
     def __init__(self,
                  root_path: str = './EEG_feature',
@@ -146,7 +139,6 @@ def process_record(file: Any = None,
                    num_channel: int = 62,
                    before_trial: Union[None, Callable] = None,
                    offline_transform: Union[None, Callable] = None,
-                   after_trial: Union[None, Callable] = None,
                    **kwargs):
         file_name = file
 
diff --git a/torcheeg/datasets/module/emotion_recognition/seed.py b/torcheeg/datasets/module/emotion_recognition/seed.py
index 46fe575..43b91c9 100644
--- a/torcheeg/datasets/module/emotion_recognition/seed.py
+++ b/torcheeg/datasets/module/emotion_recognition/seed.py
@@ -2,9 +2,6 @@
 from typing import Any, Callable, Dict, Tuple, Union
 
 import scipy.io as scio
-
-from ...constants.emotion_recognition.seed import (SEED_ADJACENCY_MATRIX,
-                                                   SEED_CHANNEL_LOCATION_DICT)
 from ..base_dataset import BaseDataset
 
 
@@ -122,8 +119,6 @@ class SEEDDataset(BaseDataset):
         verbose (bool): Whether to display logs during processing, such as progress bars, etc. (default: :obj:`True`)
         in_memory (bool): Whether to load the entire dataset into memory. If :obj:`in_memory` is set to True, then the first time an EEG sample is read, the entire dataset is loaded into memory for subsequent retrieval. Otherwise, the dataset is stored on disk to avoid the out-of-memory problem. (default: :obj:`False`)    
     '''
-    channel_location_dict = SEED_CHANNEL_LOCATION_DICT
-    adjacency_matrix = SEED_ADJACENCY_MATRIX
 
     def __init__(self,
                  root_path: str = './Preprocessed_EEG',
@@ -175,7 +170,6 @@ def process_record(file: Any = None,
                    num_channel: int = 62,
                    before_trial: Union[None, Callable] = None,
                    offline_transform: Union[None, Callable] = None,
-                   after_trial: Union[None, Callable] = None,
                    **kwargs):
         file_name = file
 
diff --git a/torcheeg/datasets/module/emotion_recognition/seed_feature.py b/torcheeg/datasets/module/emotion_recognition/seed_feature.py
index 78fd5ac..69ee406 100644
--- a/torcheeg/datasets/module/emotion_recognition/seed_feature.py
+++ b/torcheeg/datasets/module/emotion_recognition/seed_feature.py
@@ -4,11 +4,6 @@
 
 import numpy as np
 import scipy.io as scio
-
-from torcheeg.io import EEGSignalIO, MetaInfoIO
-
-from ...constants.emotion_recognition.seed import (SEED_ADJACENCY_MATRIX,
-                                                   SEED_CHANNEL_LOCATION_DICT)
 from ..base_dataset import BaseDataset
 
 
@@ -104,8 +99,6 @@ class SEEDFeatureDataset(BaseDataset):
         verbose (bool): Whether to display logs during processing, such as progress bars, etc. (default: :obj:`True`)
         in_memory (bool): Whether to load the entire dataset into memory. If :obj:`in_memory` is set to True, then the first time an EEG sample is read, the entire dataset is loaded into memory for subsequent retrieval. Otherwise, the dataset is stored on disk to avoid the out-of-memory problem. (default: :obj:`False`)    
     '''
-    channel_location_dict = SEED_CHANNEL_LOCATION_DICT
-    adjacency_matrix = SEED_ADJACENCY_MATRIX
 
     def __init__(self,
                  root_path: str = './ExtractedFeatures',
@@ -154,7 +147,6 @@ def process_record(file: Any = None,
                        num_channel: int = 62,
                        before_trial: Union[None, Callable] = None,
                        offline_transform: Union[None, Callable] = None,
-                       after_trial: Union[None, Callable] = None,
                        **kwargs):
         file_name, session_id = file
 
diff --git a/torcheeg/datasets/module/emotion_recognition/seed_iv.py b/torcheeg/datasets/module/emotion_recognition/seed_iv.py
index b6098bb..df043d1 100644
--- a/torcheeg/datasets/module/emotion_recognition/seed_iv.py
+++ b/torcheeg/datasets/module/emotion_recognition/seed_iv.py
@@ -3,11 +3,6 @@
 from typing import Any, Callable, Dict, Tuple, Union
 
 import scipy.io as scio
-
-from torcheeg.io import EEGSignalIO, MetaInfoIO
-
-from ...constants.emotion_recognition.seed_iv import (
-    SEED_IV_ADJACENCY_MATRIX, SEED_IV_CHANNEL_LOCATION_DICT)
 from ..base_dataset import BaseDataset
 
 
@@ -116,8 +111,6 @@ class SEEDIVDataset(BaseDataset):
         verbose (bool): Whether to display logs during processing, such as progress bars, etc. (default: :obj:`True`)
         in_memory (bool): Whether to load the entire dataset into memory. If :obj:`in_memory` is set to True, then the first time an EEG sample is read, the entire dataset is loaded into memory for subsequent retrieval. Otherwise, the dataset is stored on disk to avoid the out-of-memory problem. (default: :obj:`False`)    
     '''
-    channel_location_dict = SEED_IV_ADJACENCY_MATRIX
-    adjacency_matrix = SEED_IV_CHANNEL_LOCATION_DICT
 
     def __init__(self,
                  root_path: str = './eeg_raw_data',
@@ -168,7 +161,6 @@ def process_record(file: Any = None,
                    num_channel: int = 62,
                    before_trial: Union[None, Callable] = None,
                    offline_transform: Union[None, Callable] = None,
-                   after_trial: Union[None, Callable] = None,
                    **kwargs):
         file_path = file  # an element from file name list
 
diff --git a/torcheeg/datasets/module/emotion_recognition/seed_iv_feature.py b/torcheeg/datasets/module/emotion_recognition/seed_iv_feature.py
index e1d6a2e..5d97492 100644
--- a/torcheeg/datasets/module/emotion_recognition/seed_iv_feature.py
+++ b/torcheeg/datasets/module/emotion_recognition/seed_iv_feature.py
@@ -4,11 +4,6 @@
 
 import numpy as np
 import scipy.io as scio
-
-from torcheeg.io import EEGSignalIO, MetaInfoIO
-
-from ...constants.emotion_recognition.seed_iv import (
-    SEED_IV_ADJACENCY_MATRIX, SEED_IV_CHANNEL_LOCATION_DICT)
 from ..base_dataset import BaseDataset
 
 
@@ -93,8 +88,6 @@ class SEEDIVFeatureDataset(BaseDataset):
         verbose (bool): Whether to display logs during processing, such as progress bars, etc. (default: :obj:`True`)
         in_memory (bool): Whether to load the entire dataset into memory. If :obj:`in_memory` is set to True, then the first time an EEG sample is read, the entire dataset is loaded into memory for subsequent retrieval. Otherwise, the dataset is stored on disk to avoid the out-of-memory problem. (default: :obj:`False`)    
     '''
-    channel_location_dict = SEED_IV_CHANNEL_LOCATION_DICT
-    adjacency_matrix = SEED_IV_ADJACENCY_MATRIX
 
     def __init__(self,
                  root_path: str = './eeg_feature_smooth',
@@ -141,7 +134,6 @@ def process_record(feature: list = ['de_movingAve'],
                        num_channel: int = 62,
                        offline_transform: Union[None, Callable] = None,
                        before_trial: Union[None, Callable] = None,
-                       after_trial: Union[Callable, None] = None,
                        file: Any = None,
                        **kwargs):
         file_path = file  # an element from file name list
diff --git a/torcheeg/datasets/module/motor_imagery/__init__.py b/torcheeg/datasets/module/motor_imagery/__init__.py
new file mode 100644
index 0000000..25a2b58
--- /dev/null
+++ b/torcheeg/datasets/module/motor_imagery/__init__.py
@@ -0,0 +1,1 @@
+from .bciciv_2a import BCICIV2aDataset
\ No newline at end of file
diff --git a/torcheeg/datasets/module/motor_imagery/bciciv_2a.py b/torcheeg/datasets/module/motor_imagery/bciciv_2a.py
new file mode 100644
index 0000000..0d04ba7
--- /dev/null
+++ b/torcheeg/datasets/module/motor_imagery/bciciv_2a.py
@@ -0,0 +1,262 @@
+import os
+import logging
+from typing import Any, Callable, Dict, Tuple, Union
+import numpy as np
+import scipy.io as scio
+
+from ..base_dataset import BaseDataset
+
+log = logging.getLogger(__name__)
+
+
+class BCICIV2aDataset(BaseDataset):
+    r'''
+    A dataset for motor imagery, BCI Competition 2008 Graz data set A (BCICIV_2a). This class generates training samples and test samples according to the given parameters, and caches the generated results in a unified input and output format (IO). The relevant information of the dataset is as follows:
+
+    - Author: Tangermann et al.
+    - Year: 2012
+    - Download URL: http://bnci-horizon-2020.eu/database/data-sets
+    - Reference: Tangermann M, Müller K R, Aertsen A, et al. Review of the BCI competition IV[J]. Frontiers in neuroscience, 2012: 55.
+    - Signals: Electroencephalogram (22 channels at 250Hz) and electrocardiogram (3 channels at 250Hz) of 9 subjects. Two sessions on different days were recorded for each subject. Each session is comprised of 6 runs separated by short breaks. One run consists of 48 trials (12 for each of the four possible classes), yielding a total of 288 trials per session. For more detail, please refer to https://www.bbci.de/competition/iv/desc_2a.pdf
+    - Category: The imagination of movement of the left hand (class 1), right hand (class 2), both
+    feet (class 3), and tongue (class 4).
+
+    In order to use this dataset, the download folder :obj:`BCICIV_2a_mat` is required, containing the following files:
+
+    - http://bnci-horizon-2020.eu/database/data-sets/001-2014/A01T.mat
+    - http://bnci-horizon-2020.eu/database/data-sets/001-2014/A01E.mat
+    - ...
+    - http://bnci-horizon-2020.eu/database/data-sets/001-2014/A09E.mat
+
+    An example dataset for CNN-based methods:
+
+    .. code-block:: python
+
+        dataset = BCICIV2aDataset(io_path=f'./bciciv_2a',
+                              root_path='./BCICIV_2a_mat',
+                              online_transform=transforms.Compose([
+                                  transforms.To2d(),
+                                  transforms.ToTensor()
+                              ]),
+                              label_transform=transforms.Compose([
+                                  transforms.Select('label'),
+                                  transforms.Lambda(lambda x: x - 1)
+                              ]))
+        print(dataset[0])
+        # EEG signal (torch.Tensor[1, 22, 1750])
+        # label (int)
+
+        model = EEGNet(chunk_size=1750,
+                       num_electrodes=22,
+                       dropout=0.5,
+                       kernel_1=64,
+                       kernel_2=16,
+                       F1=8,
+                       F2=16,
+                       D=2,
+                       num_classes=2)
+
+    Args:
+        root_path (str): Downloaded data files in matlab (unzipped BCICIV_2a_mat.zip) formats (default: :obj:`'./BCICIV_2a_mat'`)
+        offset (int): The number of data points to be discarded at the beginning of each trial. (default: :obj:`0`)
+        chunk_size (int): Number of data points included in each EEG chunk as training or test samples. If set to -1, the EEG signal of a trial is used as a sample of a chunk. (default: :obj:`7 * 250`)
+        overlap (int): The number of overlapping data points between different chunks when dividing EEG chunks. (default: :obj:`0`)
+        num_channel (int): Number of channels used, of which the first 14 channels are EEG signals. (default: :obj:`22`)
+        skip_trial_with_artifacts (bool): Whether to skip trials with artifacts. (default: :obj:`False`)
+        online_transform (Callable, optional): The transformation of the EEG signals and baseline EEG signals. The input is a :obj:`np.ndarray`, and the ouput is used as the first and second value of each element in the dataset. (default: :obj:`None`)
+        offline_transform (Callable, optional): The usage is the same as :obj:`online_transform`, but executed before generating IO intermediate results. (default: :obj:`None`)
+        label_transform (Callable, optional): The transformation of the label. The input is an information dictionary, and the ouput is used as the third value of each element in the dataset. (default: :obj:`None`)
+        before_trial (Callable, optional): The hook performed on the trial to which the sample belongs. It is performed before the offline transformation and thus typically used to implement context-dependent sample transformations, such as moving averages, etc. The input of this hook function is a 2D EEG signal with shape (number of electrodes, number of data points), whose ideal output shape is also (number of electrodes, number of data points).
+        after_trial (Callable, optional): The hook performed on the trial to which the sample belongs. It is performed after the offline transformation and thus typically used to implement context-dependent sample transformations, such as moving averages, etc. The input and output of this hook function should be a sequence of dictionaries representing a sequence of EEG samples. Each dictionary contains two key-value pairs, indexed by :obj:`eeg` (the EEG signal matrix) and :obj:`key` (the index in the database) respectively.
+        io_path (str): The path to generated unified data IO, cached as an intermediate result. (default: :obj:`./io/bciciv_2a`)
+        io_size (int): Maximum size database may grow to; used to size the memory mapping. If database grows larger than ``map_size``, an exception will be raised and the user must close and reopen. (default: :obj:`10485760`)
+        io_mode (str): Storage mode of EEG signal. When io_mode is set to :obj:`lmdb`, TorchEEG provides an efficient database (LMDB) for storing EEG signals. LMDB may not perform well on limited operating systems, where a file system based EEG signal storage is also provided. When io_mode is set to :obj:`pickle`, pickle-based persistence files are used. (default: :obj:`lmdb`)
+        num_worker (int): Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process. (default: :obj:`0`)
+        verbose (bool): Whether to display logs during processing, such as progress bars, etc. (default: :obj:`True`)
+        in_memory (bool): Whether to load the entire dataset into memory. If :obj:`in_memory` is set to True, then the first time an EEG sample is read, the entire dataset is loaded into memory for subsequent retrieval. Otherwise, the dataset is stored on disk to avoid the out-of-memory problem. (default: :obj:`False`)
+    '''
+
+    def __init__(self,
+                 root_path: str = './BCICIV_2a_mat',
+                 offset: int = 0,
+                 chunk_size: int = 7 * 250,
+                 overlap: int = 0,
+                 num_channel: int = 22,
+                 skip_trial_with_artifacts: bool = False,
+                 online_transform: Union[None, Callable] = None,
+                 offline_transform: Union[None, Callable] = None,
+                 label_transform: Union[None, Callable] = None,
+                 before_trial: Union[None, Callable] = None,
+                 after_trial: Union[Callable, None] = None,
+                 after_session: Union[Callable, None] = None,
+                 after_subject: Union[Callable, None] = None,
+                 io_path: str = './io/bciciv_2a',
+                 io_size: int = 10485760,
+                 io_mode: str = 'lmdb',
+                 num_worker: int = 0,
+                 verbose: bool = True,
+                 in_memory: bool = False):
+        # pass all arguments to super class
+        params = {
+            'root_path': root_path,
+            'offset': offset,
+            'chunk_size': chunk_size,
+            'overlap': overlap,
+            'num_channel': num_channel,
+            'skip_trial_with_artifacts': skip_trial_with_artifacts,
+            'online_transform': online_transform,
+            'offline_transform': offline_transform,
+            'label_transform': label_transform,
+            'before_trial': before_trial,
+            'after_trial': after_trial,
+            'after_session': after_session,
+            'after_subject': after_subject,
+            'io_path': io_path,
+            'io_size': io_size,
+            'io_mode': io_mode,
+            'num_worker': num_worker,
+            'verbose': verbose,
+            'in_memory': in_memory
+        }
+        super().__init__(**params)
+        # save all arguments to __dict__
+        self.__dict__.update(params)
+
+    @staticmethod
+    def process_record(file: Any = None,
+                       offset: int = 0,
+                       chunk_size: int = 128,
+                       overlap: int = 0,
+                       num_channel: int = 14,
+                       skip_trial_with_artifacts: bool = False,
+                       before_trial: Union[None, Callable] = None,
+                       offline_transform: Union[None, Callable] = None,
+                       **kwargs):
+
+        if chunk_size <= 0:
+            dynamic_chunk_size = 7 * 250
+        else:
+            dynamic_chunk_size = int(chunk_size)
+
+        # get file name without extension
+        file_name = os.path.splitext(os.path.basename(file))[0]
+        # the last letter of the file name is the session, the rest is the subject
+        subject = file_name[:-1]
+        session = file_name[-1]
+        write_pointer = 0
+
+        a_data = scio.loadmat(file)['data']
+
+        for run_id in range(0, a_data.size):
+            # a_data: (1, 9) struct, 1-3: 25 channel EOG test (eyes open, eyes closed, movement), 4-9: 6 runs
+
+            a_data1 = a_data[0, run_id]
+            a_data2 = [a_data1[0, 0]]
+            a_data3 = a_data2[0]
+            a_X = a_data3[0]
+            a_trial = a_data3[1]
+            a_y = a_data3[2]
+            a_artifacts = a_data3[5]
+
+            a_X = np.transpose(a_X)  # to channel number, data point number
+            if before_trial:
+                a_X = before_trial(a_X)
+
+            # for EOG test, a_trial is []
+            for trial_id in range(0, a_trial.size):
+                trial_meta_info = {
+                    'subject_id': subject,
+                    'trial_id': trial_id,
+                    'session': session,
+                    'subject': subject,
+                    'run': run_id
+                }
+
+                if (a_artifacts[trial_id] != 0 and skip_trial_with_artifacts):
+                    continue
+
+                start_at = int(a_trial[trial_id] + offset)
+                end_at = start_at + dynamic_chunk_size
+                step = dynamic_chunk_size - overlap
+
+                if trial_id < a_trial.size - 1:
+                    trial_end_at = int(a_trial[trial_id + 1])
+                else:
+                    trial_end_at = a_X.shape[1]
+
+                # print(start_at)
+                # print(end_at)
+                while end_at <= trial_end_at:
+                    clip_id = f'{file_name}_{write_pointer}'
+
+                    record_info = {
+                        'start_at': start_at,
+                        'end_at': end_at,
+                        'clip_id': clip_id
+                    }
+                    record_info.update(trial_meta_info)
+
+                    t_eeg = a_X[:num_channel, start_at:end_at]
+                    if offline_transform is not None:
+                        t = offline_transform(eeg=t_eeg)
+                        t_eeg = t['eeg']
+
+                    record_info['label'] = int(a_y[trial_id])
+
+                    yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
+                    # print(t_eeg.shape)
+
+                    write_pointer += 1
+
+                    start_at = start_at + step
+                    end_at = start_at + dynamic_chunk_size
+
+    def set_records(self, root_path: str = './BCICIV_2a_mat', **kwargs):
+        file_name_list = os.listdir(root_path)
+        file_path_list = [
+            os.path.join(root_path, file) for file in file_name_list
+        ]
+
+        return file_path_list
+
+    def __getitem__(self, index: int) -> Tuple[any, any, int, int, int]:
+        info = self.read_info(index)
+
+        eeg_index = str(info['clip_id'])
+        eeg_record = str(info['_record_id'])
+        eeg = self.read_eeg(eeg_record, eeg_index)
+
+        signal = eeg
+        label = info
+
+        if self.online_transform:
+            signal = self.online_transform(eeg=eeg)['eeg']
+
+        if self.label_transform:
+            label = self.label_transform(y=info)['y']
+
+        return signal, label
+
+    @property
+    def repr_body(self) -> Dict:
+        return dict(
+            super().repr_body, **{
+                'root_path': self.root_path,
+                'chunk_size': self.chunk_size,
+                'overlap': self.overlap,
+                'num_channel': self.num_channel,
+                'skip_trial_with_artifacts': self.skip_trial_with_artifacts,
+                'online_transform': self.online_transform,
+                'offline_transform': self.offline_transform,
+                'label_transform': self.label_transform,
+                'before_trial': self.before_trial,
+                'after_trial': self.after_trial,
+                'after_session': self.after_session,
+                'after_subject': self.after_subject,
+                'io_path': self.io_path,
+                'io_size': self.io_size,
+                'io_mode': self.io_mode,
+                'num_worker': self.num_worker,
+                'verbose': self.verbose,
+                'in_memory': self.in_memory
+            })
diff --git a/torcheeg/datasets/module/personal_identification/m3cv.py b/torcheeg/datasets/module/personal_identification/m3cv.py
index fb4566d..f7167b0 100644
--- a/torcheeg/datasets/module/personal_identification/m3cv.py
+++ b/torcheeg/datasets/module/personal_identification/m3cv.py
@@ -4,8 +4,6 @@
 import pandas as pd
 import scipy.io as scio
 
-from torcheeg.io import EEGSignalIO, MetaInfoIO
-
 from ..base_dataset import BaseDataset
 
 
@@ -171,7 +169,6 @@ def process_record(file: Any = None,
                    num_channel: int = 64,
                    offline_transform: Union[None, Callable] = None,
                    before_trial: Union[None, Callable] = None,
-                   after_trial: Union[None, Callable] = None,
                    **kwargs):
         start_idx, end_idx = file
 
@@ -231,24 +228,11 @@ def process_record(file: Any = None,
                     'clip_id': clip_id
                 }
                 record_info.update(epoch_meta_info)
-                if after_trial:
-                    trial_queue.append({
-                        'eeg': t_eeg,
-                        'key': clip_id,
-                        'info': record_info
-                    })
-                else:
-                    yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
+                yield {'eeg': t_eeg, 'key': clip_id, 'info': record_info}
 
                 start_at = start_at + step
                 end_at = start_at + dynamic_chunk_size
 
-            if len(trial_queue) and after_trial:
-                trial_queue = after_trial(trial_queue)
-                for obj in trial_queue:
-                    assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
-                    yield obj
-
     def set_records(self, root_path: str = './aistudio',
                    subset: str = 'Enrollment',
                    **kwargs):
diff --git a/torcheeg/datasets/module/sspev/tsu_benchmark.py b/torcheeg/datasets/module/sspev/tsu_benchmark.py
index 30abc0c..bf7a48e 100644
--- a/torcheeg/datasets/module/sspev/tsu_benchmark.py
+++ b/torcheeg/datasets/module/sspev/tsu_benchmark.py
@@ -3,12 +3,6 @@
 from typing import Any, Callable, Dict, Tuple, Union
 
 import scipy.io as scio
-
-from torcheeg.io import EEGSignalIO, MetaInfoIO
-
-from ...constants.ssvep.tsu_benchmark import (TSUBENCHMARK_ADJACENCY_MATRIX,
-                                              TSUBENCHMARK_CHANNEL_LOCATION_DICT
-                                              )
 from ..base_dataset import BaseDataset
 
 
@@ -114,8 +108,6 @@ class TSUBenckmarkDataset(BaseDataset):
         verbose (bool): Whether to display logs during processing, such as progress bars, etc. (default: :obj:`True`)
     
     '''
-    channel_location_dict = TSUBENCHMARK_CHANNEL_LOCATION_DICT
-    adjacency_matrix = TSUBENCHMARK_ADJACENCY_MATRIX
 
     def __init__(self,
                  root_path: str = './TSUBenchmark',
@@ -163,7 +155,6 @@ def process_record(file: Any = None,
                    num_channel: int = 64,
                    offline_transform: Union[None, Callable] = None,
                    before_trial: Union[None, Callable] = None,
-                   after_trial: Union[None, Callable] = None,
                    **kwargs):
         file_name = file
 
@@ -228,14 +219,7 @@ def process_record(file: Any = None,
                         'clip_id': clip_id
                     }
                     record_info.update(block_meta_info)
-                    if after_trial:
-                        block_queue.append({
-                            'eeg': t_eeg,
-                            'key': clip_id,
-                            'info': record_info
-                        })
-                    else:
-                        yield {
+                    yield {
                             'eeg': t_eeg,
                             'key': clip_id,
                             'info': record_info
@@ -244,12 +228,6 @@ def process_record(file: Any = None,
                     start_at = start_at + step
                     end_at = start_at + dynamic_chunk_size
 
-                if len(block_queue) and after_trial:
-                    block_queue = after_trial(block_queue)
-                    for obj in block_queue:
-                        assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
-                        yield obj
-
     def set_records(self, **kwargs):
         root_path = kwargs.pop('root_path', './TSUBenchmark')  # str
 
diff --git a/torcheeg/models/cnn/__init__.py b/torcheeg/models/cnn/__init__.py
index 7f6a875..d1357b6 100644
--- a/torcheeg/models/cnn/__init__.py
+++ b/torcheeg/models/cnn/__init__.py
@@ -2,7 +2,8 @@
 from .eegnet import EEGNet
 from .fbccnn import FBCCNN
 from .fbcnet import FBCNet
+from .fbmsnet import FBMSNet,Conv2dWithConstraint,MixedConv2d,LinearWithConstraint,SamePadConv2d
 from .mtcnn import MTCNN
 from .stnet import STNet
 from .tsception import TSCeption
-from .sst_emotion_net import SSTEmotionNet
\ No newline at end of file
+from .sst_emotion_net import SSTEmotionNet
diff --git a/torcheeg/models/cnn/fbmsnet.py b/torcheeg/models/cnn/fbmsnet.py
new file mode 100644
index 0000000..25d0d07
--- /dev/null
+++ b/torcheeg/models/cnn/fbmsnet.py
@@ -0,0 +1,358 @@
+import torch.nn.functional as F
+from typing import Tuple, Optional
+import torch
+import torch.nn as nn
+from .fbcnet import VarLayer, MaxLayer, StdLayer, LogVarLayer, LinearWithConstraint, MeanLayer, swish,Conv2dWithConstraint
+
+
+## CONV_SAME_PADDING
+def _calc_same_pad(i: int, k: int, s: int, d: int):
+    return max((-(i // -s) - 1) * s + (k - 1) * d + 1 - i, 0)
+
+
+def _same_pad_arg(input_size, kernel_size, stride, dilation, **_):
+    ih, iw = input_size
+    kh, kw = kernel_size
+    pad_h = _calc_same_pad(ih, kh, stride[0], dilation[0])
+    pad_w = _calc_same_pad(iw, kw, stride[1], dilation[1])
+    return [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2]
+
+def conv2d_same(x,
+                weight: torch.Tensor,
+                bias: Optional[torch.Tensor] = None,
+                stride: Tuple[int, int] = (1, 1),
+                padding: Tuple[int, int] = (0, 0),
+                dilation: Tuple[int, int] = (1, 1),
+                groups: int = 1):
+    ih, iw = x.size()[-2:]
+    kh, kw = weight.size()[-2:]
+    pad_h = _calc_same_pad(ih, kh, stride[0], dilation[0])
+    pad_w = _calc_same_pad(iw, kw, stride[1], dilation[1])
+    x = F.pad(x,
+              [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2])
+    return F.conv2d(x, weight, bias, stride, (0, 0), dilation, groups)
+
+
+class SamePadConv2d(nn.Conv2d):
+    """ Tensorflow like 'SAME' convolution wrapper for 2D convolutions
+    """
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 dilation=1,
+                 groups=1,
+                 bias=True):
+        super(SamePadConv2d,
+              self).__init__(in_channels, out_channels, kernel_size, stride, 0,
+                             dilation, groups, bias)
+
+    def forward(self, x):
+        return conv2d_same(x, self.weight, self.bias, self.stride,
+                           self.padding, self.dilation, self.groups)
+
+
+## MIX_CONV
+def create_conv2d_pad(in_chs, out_chs, kernel_size, **kwargs):
+    padding = kwargs.pop('padding', '')
+    kwargs.setdefault('bias', False)
+    padding, is_dynamic = get_padding_value(padding, kernel_size, **kwargs)
+    if is_dynamic:
+        return SamePadConv2d(in_chs, out_chs, kernel_size, **kwargs)
+    else:
+        if isinstance(kernel_size, tuple):
+            padding = (0, padding)
+        return nn.Conv2d(in_chs,
+                         out_chs,
+                         kernel_size,
+                         padding=padding,
+                         **kwargs)
+
+
+def _is_static_pad(kernel_size, stride=1, dilation=1, **_):
+    return stride == 1 and (dilation * (kernel_size - 1)) % 2 == 0
+
+
+def _get_padding(kernel_size, stride=1, dilation=1, **_):
+    if isinstance(kernel_size, tuple):
+        kernel_size = max(kernel_size)
+    padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2
+    return padding
+
+
+def _split_channels(num_chan, num_groups):
+    split = [num_chan // num_groups for _ in range(num_groups)]
+    split[0] += num_chan - sum(split)
+    return split
+
+
+def get_padding_value(padding, kernel_size, **kwargs):
+    dynamic = False
+    if isinstance(padding, str):
+        # for any string padding, the padding will be calculated for you, one of three ways
+        padding = padding.lower()
+        if padding == 'same':
+            dynamic = True
+            padding = 0
+
+        elif padding == 'valid':
+            # 'VALID' padding, same as padding=0
+            padding = 0
+        else:
+            # Default to PyTorch style 'same'-ish symmetric padding
+            dynamic = True
+            padding = 0
+    else:
+        dynamic = True
+        padding = 0
+    return padding, dynamic
+
+
+class MixedConv2d(nn.ModuleDict):
+    """ Mixed Grouped Convolution
+    Based on MDConv and GroupedConv in MixNet impl:
+      https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mixnet/custom_layers.py
+    """
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size=3,
+                 stride=1,
+                 padding='same',
+                 dilation=1,
+                 depthwise=False,
+                 **kwargs):
+        super(MixedConv2d, self).__init__()
+
+        kernel_size = kernel_size if isinstance(kernel_size,
+                                                list) else [kernel_size]
+        num_groups = len(kernel_size)
+        in_splits = _split_channels(in_channels, num_groups)
+        out_splits = _split_channels(out_channels, num_groups)
+        self.in_channels = sum(in_splits)
+        self.out_channels = sum(out_splits)
+
+        for idx, (k, in_ch,
+                  out_ch) in enumerate(zip(kernel_size, in_splits,
+                                           out_splits)):
+            conv_groups = out_ch if depthwise else 1
+            self.add_module(
+                str(idx),
+                create_conv2d_pad(in_ch,
+                                  out_ch,
+                                  k,
+                                  stride=stride,
+                                  padding=padding,
+                                  dilation=dilation,
+                                  groups=conv_groups,
+                                  **kwargs))
+        self.splits = in_splits
+
+    def forward(self, x):
+        x_split = torch.split(x, self.splits, 1)
+        x_out = [conv(x_split[i]) for i, conv in enumerate(self.values())]
+        x = torch.cat(x_out, 1)
+        return x
+
+
+
+
+
+class FBMSNet(nn.Module):
+    r'''
+        FBMSNet, a novel multiscale temporal convolutional neural network for MI decoding tasks, employs Mixed Conv to extract multiscale temporal features which  enhance the intra-class compactness and improve the inter-class separability with the joint supervision of the center loss andcenter loss.
+
+        - Paper: FBMSNet: A Filter-Bank Multi-Scale Convolutional Neural Network for EEG-Based Motor Imagery Decoding
+        - URL: https://ieeexplore.ieee.org/document/9837422
+        - Related Project: https://github.com/Want2Vanish/FBMSNet
+
+        Below is a example to explain how to use this model. Firstly we should transform eeg signal to several nonoverlapping frequency bands by :obj:`torcheeg.transforms.BandSignal` 
+
+        .. code-block:: python
+
+            # Define 9 nonoverlapping frequency bands, each with a 4 Hz bandwidth and spanning from 4 to 40 Hz.
+            Freq_range_per_band = {'sub band1': [4, 8],
+                                'sub band2': [8, 12],
+                                'sub band3': [12, 16],
+                                'sub band4': [16, 20],
+                                'sub band5': [20, 24],
+                                'sub band6': [24, 28],
+                                'sub band7': [28, 32],
+                                'sub band8': [32, 36],
+                                'sub band9': [36, 40]}
+            dataset =BCICIV2aDataset(io_path=f'./tmp_out/bciciv2a/band_9_filters',
+                                    root_path='./BCICIV_2a_mat',
+                                    chunk_size=512,
+                                    offline_transform=transforms.BandSignal(band_dict=Freq_range_per_band,
+                                                                            sampling_rate=250),
+                                    online_transform=transforms.ToTensor(),
+                                    label_transform=transforms.Compose([
+                                    transforms.Select('label'),
+                                    transforms.Lambda(lambda x:x-1),
+                        ]))
+            data = Dataloader(dataset)
+
+            model = FBMSNet(num_classes=4,
+                            num_electrodes=22,
+                            chunk_size=512,
+                            in_channels=9 )
+
+        There are two ways to use the model. The first one, the effective way to get the prediction result but it don't output the decoded feature. 
+
+        .. code-block:: python
+
+            x,y = next(iter(data))
+            pred = model(x)
+
+        To obtain the decoded feature, use :obj:`decoder` method. If we want to obtain prediction results based on the encoded features, use :obj:`classifier` method.
+
+        .. code-block:: python
+
+            x,y = next(iter(data))
+            feature = model.decoder(x)
+            pred = model.classifier(feature)
+
+        Args:
+            num_electrodes (int): The number of electrodes. 
+            chunk_size (int): Number of data points included in each EEG chunk. 
+            in_channels (int): The number of channels of the signal corresponding to each electrode. If the original signal is used as input, in_channels is set to 1; if the original signal is split into multiple sub-bands, in_channels is set to the number of bands. (default: :obj:`9`)
+            num_classes (int): The number of classes to predict. (default: :obj:`4`)
+            stride_factor (int): The stride factor. Please make sure the chunk_size parameter is a  multiple of stride_factor parameter in order to init model successfully. (default: :obj:`4`)
+            temporal (str): The temporal layer used, with options including VarLayer, StdLayer, LogVarLayer, MeanLayer, and MaxLayer, used to compute statistics using different techniques in the temporal dimension. (default: :obj:`LogVarLayer`)
+            num_feature (int): The number of Mixed Conv output channels which can stand for various kinds of feature. (default: :obj:`36`)
+            dilatability (int): The expansion multiple of the channels after the input bands pass through spatial convolutional blocks. (default: :obj:`8`
+
+        .. automethod:: decoder
+        .. automethod:: classifier
+    '''
+
+    def __init__(self,
+                 in_channels: int,
+                 num_electrodes: int,
+                 chunk_size: int,
+                 num_classes: int = 4,
+                 stride_factor: int = 4,
+                 temporal: str = 'LogVarLayer',
+                 num_feature: int = 36,
+                 dilatability: int = 8):
+
+        super(FBMSNet, self).__init__()
+
+        self.num_classes = num_classes
+        self.in_channels = in_channels
+        self.num_electrodes = num_electrodes
+        self.chunk_size = chunk_size
+        self.stride_factor = stride_factor
+
+        try:
+            self.mixConv2d = nn.Sequential(
+                MixedConv2d(in_channels=in_channels,
+                            out_channels=num_feature,
+                            kernel_size=[(1, 15), (1, 31), (1, 63), (1, 125)],
+                            stride=1,
+                            padding='',
+                            dilation=1,
+                            depthwise=False),
+                nn.BatchNorm2d(num_feature),
+            )
+            self.scb = self.SCB(in_chan=num_feature,
+                                out_chan=num_feature * dilatability,
+                                num_electrodes=int(num_electrodes))
+
+            # Formulate the temporal agreegator
+            if temporal == 'VarLayer':
+                self.temporal_layer = VarLayer(dim=3)
+            elif temporal == 'StdLayer':
+                self.temporal_layer = StdLayer(dim=3)
+            elif temporal == 'LogVarLayer':
+                self.temporal_layer = LogVarLayer(dim=3)
+            elif temporal == 'MeanLayer':
+                self.temporal_layer = MeanLayer(dim=3)
+            elif temporal == 'MaxLayer':
+                self.temporal_layer = MaxLayer(dim=3)
+            else:
+                raise NotImplementedError
+
+            self.center_dim = self.feature_dim(in_channels, num_electrodes,
+                                               chunk_size)[-1]
+
+            self.fc = self.LastBlock(self.center_dim, num_classes)
+        except:
+            raise Exception(
+                "Model init failed: The Chunksize must be a  multiple of stride_factor.Please modify values of stride_factor or chunk_size."
+            )
+
+    def SCB(self,
+            in_chan,
+            out_chan,
+            num_electrodes,
+            weight_norm=True,
+            *args,
+            **kwargs):
+        return nn.Sequential(
+            Conv2dWithConstraint(in_chan,
+                                 out_chan, (num_electrodes, 1),
+                                 groups=in_chan,
+                                 max_norm=2,
+                                 weight_norm=weight_norm,
+                                 padding=0), nn.BatchNorm2d(out_chan), swish())
+
+    def LastBlock(self, inF, outF, weight_norm=True, *args, **kwargs):
+        return nn.Sequential(
+            LinearWithConstraint(inF,
+                                 outF,
+                                 max_norm=0.5,
+                                 weight_norm=weight_norm,
+                                 *args,
+                                 **kwargs), nn.LogSoftmax(dim=1))
+
+    def decoder(self, x):
+        r'''
+        Args:
+            x (torch.Tensor): EEG signal representation, the ideal input shape is :obj:`[n, in_channel, num_electrodes, chunk_size ]`. Here, :obj:`n` corresponds to the batch size, :obj:`in_channels` corresponds to the number of sub bands.
+
+        Returns:
+           torch.Tensor[size of batch, length of deep feature code]: The extracted deep features.
+        '''
+        x = self.mixConv2d(x)
+        x = self.scb(x)
+        x = x.reshape([
+            *x.shape[0:2], self.stride_factor,
+            int(x.shape[3] / self.stride_factor)
+        ])
+        x = self.temporal_layer(x)
+        return torch.flatten(x, start_dim=1)
+
+    def classifier(self, feature):
+        r'''
+        With feature which is ouput by decoder inputed,the classifier ouput the predicted probability that the samples belong to the classes. 
+
+        Args:
+            feature (torch.Tensor): The extracted deep features. The ideal input shape is :obj:`[batch size,1152]`where feature dim is fixed as :obj:`1152`.
+        Returns:
+           torch.Tensor[size of batch, num_classes]: The predicted probability that the samples belong to the classes.
+        '''
+        return self.fc(feature)
+
+    def forward(self, x):
+        r'''
+        Args:
+            x (torch.Tensor): EEG signal representation, the ideal input shape is :obj:`[n, in_channel, num_electrodes, chunk_size ]`. Here, :obj:`n` corresponds to the batch size
+
+        Returns:
+            torch.Tensor[size of batch,number of classes]: The predicted probability that the samples belong to the classes.
+        '''
+        f = self.decoder(x)
+        return self.classifier(f)
+
+    def feature_dim(self, in_channels, num_electrodes, chunk_size):
+        data = torch.ones((1, in_channels, num_electrodes, chunk_size))
+        x = self.mixConv2d(data)
+        x = self.scb(x)
+        x = x.reshape([*x.shape[0:2], self.stride_factor, -1])
+        x = self.temporal_layer(x)
+        x = torch.flatten(x, start_dim=1)
+        return x.size()
diff --git a/torcheeg/models/gan/__init__.py b/torcheeg/models/gan/__init__.py
index fa19792..b31f43d 100644
--- a/torcheeg/models/gan/__init__.py
+++ b/torcheeg/models/gan/__init__.py
@@ -1,2 +1,3 @@
 from .bcgan import *
-from .bgan import *
\ No newline at end of file
+from .bgan import *
+from .eegfusenet import *
\ No newline at end of file
diff --git a/torcheeg/models/gan/eegfusenet.py b/torcheeg/models/gan/eegfusenet.py
new file mode 100644
index 0000000..94007a2
--- /dev/null
+++ b/torcheeg/models/gan/eegfusenet.py
@@ -0,0 +1,230 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+
+class EEGfuseNet(nn.Module):
+    r'''
+    EEGFuseNet: A hybrid unsupervised network which can fuse high-dimensional EEG to obtain deep feature characterization and generate similar signals. For more details, please refer to the following information.
+
+    - Paper: Z. Liang, R. Zhou, L. Zhang, L. Li, G. Huang, Z. Zhang, and S. Ishii, EEGFuseNet: Hybrid Unsupervised Deep Feature Characterization and Fusion for High-Dimensional EEG With an #Application to Emotion Recognition, IEEE Transactions on Neural Systems and Rehabilitation Engineering, 29, pp. 1913-1925, 2021.
+    - URL: https://github.com/KAZABANA/EEGfusenet
+
+    .. code-block:: python
+        g_model=EEGfuseNet(num_electrodes=20,hid_channels_gru=16,num_layers_gru=1,hid_channels_cnn=1,chunk_size=128)
+        X = torch.rand(2,1,20,128)
+        fake_X,deep_feature=g_model(X)
+
+    Args:
+        in_channels (int): The number of channels of the signal corresponding to each electrode. If the original signal is used as input, in_channels is set to 1; if the original signal is split into multiple sub-bands, in_channels is set to the number of bands. (default: :obj:`1`)
+        num_electrodes (int): The number of electordes in input data. (default: :obj:`32`)
+        hid_channels_gru (int): The number of hidden nodes in BI-GRU (default: :obj:`16`)
+        num_layers_gru (int): The number of layers of BI-GRU. (default: :obj:`1`)
+        hid_channels_cnn (int): The number of filters of CNN based encoder. (default: :obj:`1`)
+        chunk_size (int): The number of data points included in each EEG chunk.the size of the input EEG signal is( batch size × Channel × Time) (default: :obj:`384`)
+    '''
+    def __init__(self,
+                    in_channels: int = 1,
+                    num_electrodes: int = 32,
+                    hid_channels_gru: int = 16,
+                    num_layers_gru: int = 1,
+                    hid_channels_cnn: int = 1,
+                    chunk_size:int = 384 ):  
+        super(EEGfuseNet, self).__init__()
+        
+        
+        self.hid_channels_cnn = hid_channels_cnn
+        self.hid_channels_gru =  hid_channels_gru
+        self.length = chunk_size/32
+
+        self.conv1 = nn.Conv2d(in_channels, 16*hid_channels_cnn, (1, int(chunk_size/2+1)),stride = 1, padding = (0,int(chunk_size/4)))
+        self.batchNorm1 = nn.BatchNorm2d(16*hid_channels_cnn, False)
+        
+        self.depthwiseconv2 = nn.Conv2d(16*hid_channels_cnn,32*hid_channels_cnn,(num_electrodes,1),padding = 0)
+        self.batchNorm2 = nn.BatchNorm2d(32*hid_channels_cnn,False)
+        self.pooling1 = nn.MaxPool2d((1,4),return_indices =True,ceil_mode=True)
+    
+        self.separa1conv3 = nn.Conv2d(32*hid_channels_cnn,32*hid_channels_cnn,(1, int(chunk_size/8+1)),
+                                        stride=1,
+                                        padding=(0,int(chunk_size/16)),
+                                        groups=int(32*hid_channels_cnn)) 
+        self.separa2conv4 = nn.Conv2d(32*hid_channels_cnn,16*hid_channels_cnn,1) 
+        self.batchNorm3 = nn.BatchNorm2d(16*hid_channels_cnn,False)
+        self.pooling2 = nn.MaxPool2d((1,8), 
+                                        return_indices =True,
+                                        ceil_mode=True)
+        self.dropout1 = nn.Dropout(p=0.25)
+        self.dropout2 = nn.Dropout(p=0.25)
+        self.dropout3 = nn.Dropout(p=0.25)
+        self.dropout4 = nn.Dropout(p=0.25)
+        self.fc1 = nn.Linear(16*hid_channels_cnn,16*hid_channels_cnn)
+        
+        self.fc2 = nn.Linear(hid_channels_gru*2*hid_channels_cnn,
+                                hid_channels_gru*hid_channels_cnn)
+        
+        self.fc3 = nn.Linear(hid_channels_gru*hid_channels_cnn,
+                                hid_channels_gru*2*hid_channels_cnn)
+        
+        self.fc4 = nn.Linear(2*16*hid_channels_cnn,
+                                16*hid_channels_cnn)
+
+        self.gru_en = nn.GRU(16*hid_channels_cnn,hid_channels_gru*hid_channels_cnn,num_layers_gru,
+                                batch_first=True,
+                                bidirectional=True)
+        
+        self.gru_de = nn.GRU(2*hid_channels_gru*hid_channels_cnn,16*hid_channels_cnn,num_layers_gru,
+                                batch_first=True,
+                                bidirectional=True)
+        
+        self.lstm = nn.LSTM(16*hid_channels_cnn,hid_channels_gru*hid_channels_cnn,num_layers_gru,
+                            batch_first=True,
+                            bidirectional=True)
+
+        self.unpooling2 = nn.MaxUnpool2d((1, 8))
+        self.batchnorm4 = nn.BatchNorm2d(32*hid_channels_cnn,False)
+        self.desepara2conv4 =  nn.ConvTranspose2d(16*hid_channels_cnn,32*hid_channels_cnn,1)
+        self.desepara1conv3 =  nn.ConvTranspose2d( 32*hid_channels_cnn, 32*hid_channels_cnn, (1, int(chunk_size/8+1)),
+                                                    stride=1,
+                                                    padding=(0,int(chunk_size/16)),
+                                                    groups= 32*hid_channels_cnn )
+    
+        self.unpooling1 = nn.MaxUnpool2d((1, 4))
+        self.batchnorm5 = nn.BatchNorm2d(16*hid_channels_cnn, False)#
+        self.dedepthsepara1conv3 = nn.ConvTranspose2d(32*hid_channels_cnn,16*hid_channels_cnn, 
+                                                      (num_electrodes, 1), 
+                                                      stride= 1, 
+                                                      padding=0)
+        self.deconv1 = nn.ConvTranspose2d(16*hid_channels_cnn, 1, (1, int(chunk_size/2+1)), 
+                                            stride = 1, 
+                                            padding =(0,int(chunk_size/4)))
+    def forward(self,x):
+        r'''
+        Args:
+            x (torch.Tensor): EEG signal representation, the ideal input shape is :obj:`[n, 32, 384]`. Here, :obj:`n` corresponds to the batch size, :obj:`32` corresponds to :obj:`num_electrodes`, and :obj:`384` corresponds to :obj:`chunk_size`.
+
+        Returns:
+            torch.Tensor[size of batch, number of electrodes, length of EEG signal]，torch.Tensor[size of batch, length of deep feature code]: The first value is generated EEG signals. The second value is batch of extracted deep features,which used in the unsupervised EEG decoding, can represent the entire input EEG signals \ cross time points covering not only the EEG characteristics but also the EEG characteristics in the sequential information.
+        '''
+
+        x = self.conv1(x)
+        x = self.batchNorm1(x)
+        
+        x = self.depthwiseconv2(x)
+        x = self.batchNorm2(x)       
+        x = F.elu(x)
+        end_dim1 = x.shape[-1]
+        x,idx1 = self.pooling1(x)  
+        x = self.dropout1(x)
+        x = self.separa1conv3(x)     
+        x = self.separa2conv4(x) 
+
+        x = self.batchNorm3(x)
+        x = F.elu(x)
+        end_dim2 = x.shape[-1]
+        x,idx2 = self.pooling2(x)
+
+        x = x.permute(0,3,2,1)
+        x = x[:,:,-1,:,]
+        x = self.fc1(x)
+        x = F.elu(x)
+        out,_ = self.gru_en(x)
+        x = out
+        x = self.fc2(x)
+        code = x.reshape((x.shape[0],-1)) 
+        x = self.fc3(x)
+        x,_ = self.gru_de(x)
+    
+        x = self.fc4(x)
+        x = F.elu(x)
+        x = x.reshape((x.shape[0],x.shape[1],1,x.shape[2]))
+        x = x.permute(0,3,2,1)
+        
+        x = self.unpooling2(x, idx2)
+        x = x[:,:,:,:end_dim2]
+
+        x = self.desepara2conv4(x)
+        x = self.desepara1conv3(x)
+        x = self.batchnorm4(x)
+        x = self.dropout4(x)
+        
+        x = F.elu(x)   
+        x = self.unpooling1(x, idx1)  
+        x = x[:,:,:,:end_dim1]    
+        x = self.dedepthsepara1conv3(x)
+        x = self.batchnorm5(x)
+
+        x = self.deconv1(x)
+        return x,code
+
+class EFDiscriminator(nn.Module):
+    r'''
+    EFDiscriminator: the discriminator that comes with EEGFuseNet is to distinguish whether the input EEG signals is a fake one generated by the eegfusenet or a real one collected from human brain.
+
+    - Paper: Z. Liang, R. Zhou, L. Zhang, L. Li, G. Huang, Z. Zhang, and S. Ishii, EEGFuseNet: Hybrid Unsupervised Deep Feature Characterization and Fusion for High-Dimensional EEG With an #Application to Emotion Recognition, IEEE Transactions on Neural Systems and Rehabilitation Engineering, 29, pp. 1913-1925, 2021.
+    - URL: https://github.com/KAZABANA/EEGfusenet
+
+    .. code-block:: python
+        g_model=EEGfuseNet(num_electrodes=20,hid_channels_gru=16,num_layers_gru=1,hid_channels_cnn=1,chunk_size=128)
+        d_model = EFDiscriminator(num_electrodes=20,hid_channels_cnn=1,chunk_size=128)
+        X = torch.rand(128,1,20,128)
+        fake_X,deep_feature=g_model(X)
+        p_real,p_fake = d_model(X),d_model(fake_X)
+
+    Args:
+        in_channels (int): The number of channels of the signal corresponding to each electrode. If the original signal is used as input, in_channels is set to 1; if the original signal is split into multiple sub-bands, in_channels is set to the number of bands. (default: :obj:`1`)
+        num_electrodes (int): The number of electrodes. (default: :obj:`32`)
+        hid_channels_cnn (int): The number of filters in CNN based encoder. (default: :obj:`1`)
+        chunk_size (int): Number of data points included in each EEG chunk. (default: :obj:`384`)
+    '''
+    def __init__(self,
+                    in_channels: int = 1,
+                    num_electrodes:int = 32,
+                    hid_channels_cnn:int = 1,
+                    chunk_size:int = 384):
+        super(EFDiscriminator, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, 8*hid_channels_cnn, (1, int(chunk_size/2+1)),
+                                stride = 1,
+                                padding = 'same')
+        self.batchNorm1 =nn.BatchNorm2d(8*hid_channels_cnn, False)
+        self.length=chunk_size/32
+
+       
+        self.depthwiseconv2 = nn.Conv2d(8* hid_channels_cnn, 16*hid_channels_cnn ,(num_electrodes,1), padding = 0)  
+        self.batchNorm2 = nn.BatchNorm2d( 16*hid_channels_cnn, False )
+        self.pooling1 = nn.MaxPool2d((1,4),
+                                        return_indices=False ,
+                                        ceil_mode = True)
+
+        self.separa1conv3 = nn.Conv2d(16*hid_channels_cnn, 16*hid_channels_cnn,
+                                        (1, int(chunk_size/8+1)),
+                                        stride=1,padding='same',
+                                        groups=16*hid_channels_cnn )
+
+        self.separa2conv4 = nn.Conv2d(16*hid_channels_cnn,8*hid_channels_cnn,1) 
+        self.batchNorm3 = nn.BatchNorm2d(8*hid_channels_cnn,False)
+        self.pooling2 = nn.MaxPool2d((1,8), return_indices = False)
+        self.fc1 = nn.Linear(int(self.length)*8*hid_channels_cnn,1)
+    
+    def forward(self,x):
+        r'''
+        Args:
+            x (torch.Tensor): EEG signal representation or the fake generated EEGsignal, the size of the input EEG signal is( batch size × Channel × Time) whose ideal input shape is :obj:`[n, 32, 384]`. Here, :obj:`n` corresponds to the batch size, :obj:`32` corresponds to :obj:`num_electrodes`, and :obj:`384` corresponds to :obj:`chunk_size`.
+
+        Returns:
+            torch.Tensor[size of batch, 1]: The possibilities that model judging the corresponding input signals is real. 
+        '''        
+        x = self.conv1(x)
+        x = self.batchNorm1(x)
+        x = self.depthwiseconv2(x)
+        x = self.batchNorm2(x)
+        x = F.elu(x)
+        x = self.pooling1(x)
+        x = self.separa1conv3(x)
+        x = self.separa2conv4(x)
+        x = self.batchNorm3(x)
+        x = F.elu(x)
+        x = self.pooling2(x)
+        x = x.reshape((x.shape[0],-1))         
+        x = self.fc1(x)
+        x=torch.sigmoid(x)
+        return x
+
diff --git a/torcheeg/models/transformer/__init__.py b/torcheeg/models/transformer/__init__.py
index 1ccf37e..7b5ddf6 100644
--- a/torcheeg/models/transformer/__init__.py
+++ b/torcheeg/models/transformer/__init__.py
@@ -2,4 +2,5 @@
 from .simple_vit import SimpleViT
 from .vanilla_transformer import VanillaTransformer
 from .vit import ViT
-from .conformer import Conformer
\ No newline at end of file
+from .conformer import Conformer
+from .atcnet import ATCNet
\ No newline at end of file
diff --git a/torcheeg/models/transformer/atcnet.py b/torcheeg/models/transformer/atcnet.py
new file mode 100644
index 0000000..bed19a3
--- /dev/null
+++ b/torcheeg/models/transformer/atcnet.py
@@ -0,0 +1,170 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+
+class ATCNet(nn.Module):
+    r'''
+    ATCNet: An attention-based temporal convolutional network forEEG-based motor imagery classiﬁcation.For more details ,please refer to the following information:
+
+    - Paper: H. Altaheri, G. Muhammad and M. Alsulaiman, "Physics-Informed Attention Temporal Convolutional Network for EEG-Based Motor Imagery Classification," in IEEE Transactions on Industrial Informatics, vol. 19, no. 2, pp. 2249-2258, Feb. 2023, doi: 10.1109/TII.2022.3197419.
+    - URL: https://github.com/Altaheri/EEG-ATCNet
+
+    .. code-block:: python
+        model = ATCNet(in_channels=1,num_classes=4,num_windows=3,num_electrodes=22,chunk_size=128)
+        X = torch.rand(2,1,22,128) # (batch_size, in_channels, num_electrodes,chunk_size) 
+        p = model(X)
+
+    Args:
+        in_channels (int): The number of channels of the signal corresponding to each electrode. If the original signal is used as input, in_channels is set to 1; if the original signal is split into multiple sub-bands, in_channels is set to the number of bands. (default: :obj:`1`)
+        num_electrodes (int): The number of electrodes. (default: :obj:`32`)
+        num_classes (int): The number of classes to predict. (default: :obj:`4`)
+        num_windows (int): The number of sliding windows after conv block. (default: :obj:`3`)
+        num_electrodes (int): The number of electrodes if the input is EEG signal. (default: :obj:`22`)
+        conv_pool_size (int):  The size of the second average pooling layer kernel in the conv block. (default: :obj:`7`)
+        F1 (int): The channel size of the temporal feature maps in conv block. (default: :obj:`16`)
+        D (int): The number of second conv layer's filters linked to each temporal feature map in the previous layer in conv block. (default: :obj:`2`)
+        tcn_kernel_size (int): The size of conv layers kernel in the TCN block. (default: :obj:`4`)
+        tcn_depth (int): The times of TCN loop. (default: :obj:`2`)
+        chunk_size (int): The Number of data points included in each EEG chunk. (default: :obj:`1125`)
+    '''
+    def __init__(self,
+                    in_channels: int = 1,
+                    num_classes: int = 4,
+                    num_windows: int = 3,
+                    num_electrodes: int = 22,
+                    conv_pool_size: int = 7,
+                    F1: int = 16,
+                    D: int =2,
+                    tcn_kernel_size: int = 4,
+                    tcn_depth: int = 2,
+                    chunk_size: int = 1125,
+                    ):  
+        super(ATCNet, self).__init__()
+        self.in_channels = in_channels
+        self.num_classes = num_classes
+        self.num_windows = num_windows
+        self.num_electrodes = num_electrodes
+        self.pool_size = conv_pool_size
+        self.F1 = F1
+        self.D = D
+        self.tcn_kernel_size = tcn_kernel_size
+        self.tcn_depth = tcn_depth
+        self.chunk_size = chunk_size
+        F2 = F1*D
+
+        self.conv_block = nn.Sequential(
+            nn.Conv2d(in_channels,F1,(1,int(chunk_size/2+1)),stride=1,padding = 'same',bias=False),
+            nn.BatchNorm2d(F1, False),
+            nn.Conv2d(F1,F2,(num_electrodes,1),padding = 0,groups=F1),
+            nn.BatchNorm2d(F2,False),
+            nn.ELU(),
+            nn.AvgPool2d((1,8)),
+            nn.Dropout2d(0.1),
+            nn.Conv2d(F2,F2,(1,16),bias=False,padding='same'),
+            nn.BatchNorm2d(F2,False),
+            nn.ELU(),
+            nn.AvgPool2d((1,self.pool_size)),
+            nn.Dropout2d(0.1)
+        )
+        self.__build_model()
+        
+    def __build_model(self):
+        with torch.no_grad():
+            x = torch.zeros(2,self.in_channels,self.num_electrodes,self.chunk_size)
+            x = self.conv_block(x)
+            x = x[:,:,-1,:]
+            x = x.permute(0,2,1)
+            self.__chan_dim,self.__embed_dim = x.shape[1:]
+            self.win_len = self.__chan_dim - self.num_windows +1
+
+            for i in range(self.num_windows):
+                st = i 
+                end = x.shape[1]  -self.num_windows+i+1
+                x2 = x[:,st:end,:]
+
+                self.__add_msa(i)
+                x2_= self.get_submodule("msa"+str(i))(x2,x2,x2)[0]
+                self.__add_msa_drop(i) 
+                x2_ = self.get_submodule("msa_drop"+str(i))(x2)
+                x2 = torch.add(x2,x2_)
+                
+                for j in range(self.tcn_depth):
+                    self.__add_tcn((i+1)*j,x2.shape[1])
+                    out = self.get_submodule("tcn"+str( (i+1)*j ))(x2)
+                    if x2.shape[1] != out.shape[1]: 
+                        self.__add_recov(i)                   
+                        x2 = self.get_submodule("re"+str(i))(x2)
+                    x2 = torch.add(x2,out)
+                    x2 = nn.ELU()(x2) 
+                x2 = x2[:,-1,:]
+                self.__dense_dim = x2.shape[-1]
+                self.__add_dense(i)
+                x2 = self.get_submodule("dense"+str(i))(x2)
+
+   
+    def __add_msa(self,index:int):
+        
+        self.add_module('msa'+str(index),nn.MultiheadAttention(
+                                         embed_dim=self.__embed_dim,
+                                         num_heads=2,
+                                         batch_first=True))
+    def __add_msa_drop(self,index):
+        self.add_module('msa_drop'+str(index),nn.Dropout(0.3))
+
+    def __add_tcn(self,index:int,num_electrodes:int):
+        self.add_module('tcn'+str(index), 
+           nn.Sequential(
+            nn.Conv1d(num_electrodes,32,self.tcn_kernel_size,padding='same'),
+            nn.BatchNorm1d(32),
+            nn.ELU(),
+            nn.Dropout(0.3),
+            nn.Conv1d(32,32,self.tcn_kernel_size,padding = 'same'),
+            nn.BatchNorm1d(32),
+            nn.ELU(),
+            nn.Dropout(0.3) )
+        )
+
+    def __add_recov(self,index:int):
+        self.add_module('re'+str(index),nn.Conv1d(self.win_len,32,4,padding='same'))
+
+    def __add_dense(self, index:int):
+        self.add_module('dense'+str(index),nn.Linear(self.__dense_dim,self.num_classes))
+
+    def forward(self,x):
+        r'''
+        Args:
+            x (torch.Tensor): EEG signal representation, the ideal input shape is :obj:`[n, 22, 1125]`. Here, :obj:`n` corresponds to the batch size, :obj:`22` corresponds to :obj:`num_electrodes`, and :obj:`1125` corresponds to :obj:`chunk_size`.
+
+        Returns:
+            torch.Tensor[size of batch, number of classes]: The predicted probability that the samples belong to the classes.
+        '''
+        x = self.conv_block(x)
+        x = x[:,:,-1,:]
+        x = x.permute(0,2,1)
+
+        
+        for i in range(self.num_windows):
+            st = i 
+            end = x.shape[1] -self.num_windows+i+1
+            x2 = x[:,st:end,:]
+            x2_= self.get_submodule("msa"+str(i))(x2,x2,x2)[0] 
+            x2_ = self.get_submodule("msa_drop"+str(i))(x2)
+            x2 = torch.add(x2,x2_)
+            
+
+            for j in range(self.tcn_depth):
+               out = self.get_submodule("tcn"+str( (i+1)*j ))(x2)
+               if x2.shape[1] != out.shape[1]:                 
+                    x2 = self.get_submodule("re"+str(i))(x2)
+               x2 = torch.add(x2,out)
+               x2 = nn.ELU()(x2) 
+            x2 = x2[:,-1,:]
+            x2 = self.get_submodule("dense"+str(i))(x2)
+            if i == 0:
+                sw_concat = x2
+            else:
+                sw_concat =sw_concat.add(x2)
+
+        x = sw_concat/self.num_windows
+        x = nn.Softmax(dim=1)(x)
+        return x
diff --git a/torcheeg/trainers/__init__.py b/torcheeg/trainers/__init__.py
index e42e3ed..a64504c 100644
--- a/torcheeg/trainers/__init__.py
+++ b/torcheeg/trainers/__init__.py
@@ -2,4 +2,5 @@
 from .domain_adaption import *
 from .generative import *
 from .self_supervised import *
-from .finetune import *
\ No newline at end of file
+from .finetune import *
+from .imbalance import *
\ No newline at end of file
diff --git a/torcheeg/trainers/domain_adaption/__init__.py b/torcheeg/trainers/domain_adaption/__init__.py
index 4b6e2c7..4646a8d 100644
--- a/torcheeg/trainers/domain_adaption/__init__.py
+++ b/torcheeg/trainers/domain_adaption/__init__.py
@@ -3,4 +3,5 @@
 from .dann import DANNTrainer
 from .dan import DANTrainer
 from .ada import ADATrainer
-from .jan import JANTrainer
\ No newline at end of file
+from .jan import JANTrainer
+from .center import CenterLossTrainer
\ No newline at end of file
diff --git a/torcheeg/trainers/domain_adaption/center.py b/torcheeg/trainers/domain_adaption/center.py
new file mode 100644
index 0000000..a271f8d
--- /dev/null
+++ b/torcheeg/trainers/domain_adaption/center.py
@@ -0,0 +1,386 @@
+from ..classifier import ClassifierTrainer, classification_metrics
+import torch
+import torch.nn as nn
+from torch.autograd.function import Function
+from numpy import random
+from typing import List, Tuple
+import torchmetrics
+
+
+class CentersLoss(nn.Module):
+    '''
+    ClassCentersLoss
+    '''
+    def __init__(self, num_centers, center_dim, size_average=True):
+        super(CentersLoss, self).__init__()
+        centers = random.randn(num_centers, center_dim)
+        self.centers = nn.Parameter(torch.from_numpy(centers))
+        self.ClassCentersfunc = ClassCentersFunc.apply
+        self.feat_dim = center_dim
+        self.size_average = size_average
+
+    def forward(self, label, feat):
+        batch_size = feat.size(0)
+        feat = feat.view(batch_size, -1)
+        # To check the dim of centers and features
+        if feat.size(1) != self.feat_dim:
+            raise ValueError(
+                "Center's dim: {0} should be equal to input feature's \
+                            dim: {1}".format(self.feat_dim, feat.size(1)))
+        batch_size_tensor = feat.new_empty(1).fill_(
+            batch_size if self.size_average else 1)
+        loss = self.ClassCentersfunc(feat, label, self.centers,
+                                     batch_size_tensor)
+        return loss
+
+
+class ClassCentersFunc(Function):
+
+    @staticmethod
+    def forward(ctx, feature, label, centers, batch_size):
+        ctx.save_for_backward(feature, label, centers, batch_size)
+        centers_batch = centers.index_select(0, label.long())
+        return (feature - centers_batch).pow(2).sum() / 2.0 / batch_size
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        feature, label, centers, batch_size = ctx.saved_tensors
+        centers_batch = centers.index_select(0, label.long())
+        diff = centers_batch - feature
+        # init every iteration
+        counts = centers.new_ones(centers.size(0))
+        ones = centers.new_ones(label.size(0))
+        grad_centers = centers.new_zeros(centers.size())
+
+        counts = counts.scatter_add_(0, label.long(), ones)
+        grad_centers.scatter_add_(
+            0,
+            label.unsqueeze(1).expand(feature.size()).long(), diff)
+        grad_centers = grad_centers / counts.view(-1, 1)
+        return -grad_output * diff / batch_size, None, grad_centers / batch_size, None
+
+
+
+class CenterLossTrainer(ClassifierTrainer):
+    r'''
+        A trainer trains classification model contains a decoder and a classifier. As for Center loss, it can make the output of the decoder close to the mean of decoded features within the same class. PLease refer to the following infomation to comprehend how the center loss works.
+
+        - Paper: FBMSNet: A Filter-Bank Multi-Scale Convolutional Neural Network for EEG-Based Motor Imagery Decoding
+        - URL: https://ieeexplore.ieee.org/document/9837422
+        - Related Project: https://github.com/Want2Vanish/FBMSNet
+
+        .. code-block:: python
+
+            trainer = CenterLossTrainer(decoder = decoder, 
+                                             classifier = classifier,
+                                             num_classes = your_classes,
+                                             feature_dim = your_decoded_dim)
+                                             
+            trainer.fit(train_loader, val_loader)
+            trainer.test(test_loader)
+
+        The model structure is required to contains a decoder block which generates the deep feature code and a classifier connected to the decoder to judge which class the feature code belong to.
+        Firstly, we should prepare a :obj:`decoder` model and a :obj:`classifier` model for  decoding and classifying from decoding ouput respectly. 
+        Here we take FBMSNet as example. :obj:`torcheeg.models.FBMSNet` contains decoder and classifer method already and what We need to do is just to inherit the model to define a decoder and a classifier,and then override the :obj:`forward` method . 
+
+        .. code-block:: python
+
+            from torcheeg.models import FBMSNet
+
+            class FBMSDecoder(FBMSNet):
+                def forward(self,x):
+                    return self.decoder(x)
+
+            class FBMSClassifier(FBMSNet):
+                def forward(self,x):
+                    return 
+                
+            decoder  = FBMSDecoder(num_classes=4,
+                          num_electrodes=22,
+                          chunk_size=512,
+                         in_channels=9)
+
+            classifier = FBMSClassifier(num_classes=4,
+                                num_electrodes=22,
+                                chunk_size=512,
+                                in_channels=9)
+            
+            trainer = CenterLossTrainer(decoder = decoder, 
+                                             classifier = classifier,
+                                             num_classes = 4,
+                                             feature_dim = 1152)
+        
+        Custom model is OK. Feel free to refer to  this example:
+    
+        .. code-block:: python
+
+
+            class MyDecoder(nn.Module):
+                def __init__(self):
+                    self.layer = nn.Linear(128,64)   #(input dim, decoded dim)
+
+                def forward(self,x):
+                    return self.layer(x)
+
+            class MyClassifier(nn.Module):
+                def __init__(self):
+                    self.layer = nn.Linear(64,2)      #(decoded dim, num_classes)
+
+                def forward(self,x):classifier
+                    return self.layer(x)
+                
+            decoder  = MyDecoder()
+            classifier = MyClassifier()
+
+            trainer = CenterLossTrainer(decoder = decoder, 
+                                             classifier = classifier,
+                                             num_classes = 2,
+                                             feature_dim = 64)
+        
+        Args:
+            decoder (nn.Module): The decoder which transforms eegsignal into 1D feature code.
+            classifier (nn.Module): The classifier that predict from the decoder output which class the siginals belong to.
+            feature_dim (int): The dimemsion of decoder output code whose mean values we can loosely regard as the "center". 
+            num_classes (int, optional): The number of categories in the dataset. 
+            lammda (float): The weight of the center loss in total loss. (default: :obj:`5e-4`)
+            lr (float): The learning rate. (default: :obj:`0.001`)
+            weight_decay (float): The weight decay. (default: :obj:`0.0`)
+            devices (int): The number of devices to use. (default: :obj:`1`)
+            accelerator (str): The accelerator to use. Available options are: 'cpu', 'gpu'. (default: :obj:`"cpu"`)
+            metrics (list of str): The metrics to use. Available options are: 'precision', 'recall', 'f1score', 'accuracy'. (default: :obj:`['accuracy', 'precision', 'recall', 'f1score']`)
+
+        .. automethod:: fit
+        .. automethod:: test
+    '''
+    def __init__(self,
+                 decoder,
+                 classifier,
+                 feature_dim: int,
+                 num_classes: int,
+                 lammda: float = 0.0005,
+                 lr: float = 1e-3,
+                 weight_decay: float = 0.0,
+                 devices: int = 1,
+                 accelerator: str = "cpu",
+                 metrics: List[str] = [
+                     'accuracy', 'precision', 'recall', 'f1score']
+                 ):
+        
+        super(CenterLossTrainer,self).__init__(decoder, num_classes, lr, weight_decay, devices,
+                             accelerator, metrics)
+        
+        self.decoder = decoder
+        self.classifier = classifier
+        self.center_loss = CentersLoss(num_classes,feature_dim)
+        self.lammda = lammda
+        self.automatic_optimization = False
+        self.feature_dim = feature_dim
+    
+        
+    
+    def init_metrics(self, metrics: List[str], num_classes: int) -> None:
+        # for train
+        self.train_loss = torchmetrics.MeanMetric()
+        self.center_loss_train = torchmetrics.MeanMetric()
+        self.predict_loss_train = torchmetrics.MeanMetric()
+        
+        # val
+        self.val_loss = torchmetrics.MeanMetric()
+        self.center_loss_val = torchmetrics.MeanMetric()
+        self.predict_loss_val = torchmetrics.MeanMetric()
+
+        #test
+        self.test_loss = torchmetrics.MeanMetric()
+        self.center_loss_test = torchmetrics.MeanMetric()
+        self.predict_loss_test= torchmetrics.MeanMetric()
+
+        # classification metrics for train val test
+        self.train_metrics = classification_metrics(metrics, num_classes)
+        self.val_metrics = classification_metrics(metrics, num_classes)
+        self.test_metrics = classification_metrics(metrics, num_classes)
+    
+    def __reset_metric(self, state: str = "train"):
+        if state == "train":
+            self.train_loss.reset()
+            self.center_loss_train.reset()
+            self.predict_loss_train.reset()
+            self.train_metrics.reset()
+        elif state == "val":
+            self.val_loss.reset()
+            self.center_loss_val.reset()
+            self.predict_loss_val.reset()
+            self.val_metrics.reset()
+        elif state == "test":
+            self.test_loss.reset()
+            self.center_loss_test.reset()
+            self.predict_loss_test.reset()
+            self.test_metrics.reset()
+
+
+    def training_step(self, batch, batch_idx) -> None:
+        x, y = batch
+        center_optimizer, model_optimizer = self.optimizers(True)
+
+        # zero_grad
+        center_optimizer.zero_grad()
+        model_optimizer.zero_grad()
+
+        # center loss
+        feat = self.decoder(x)
+        centerloss = self.center_loss(y, feat)
+
+        # prediction cross entropy loss
+        y_hat = self.classifier(feat)
+        pre_loss = self.ce_fn(y_hat, y)
+
+        # total loss
+        loss = self.lammda * centerloss + pre_loss
+
+        # backward
+        self.manual_backward(loss)
+  
+        # step
+        center_optimizer.step()
+        model_optimizer.step()
+  
+        # log
+        log_dict = {"train_loss" : self.train_loss(loss)}
+
+        for i, metric_value in enumerate(self.train_metrics.values()):
+            log_dict[f"train_{self.metrics[i]}"] = metric_value(y_hat, y)
+        self.log_dict(log_dict,
+                     prog_bar=True,
+                     on_epoch=False,
+                     logger=False,
+                     on_step=True)
+        
+
+
+    def on_train_epoch_end(self) -> None:
+
+        # log loss
+        log_dict = {"train_loss" : self.train_loss.compute(), }
+        # log classfication metrics
+        for i, metric_value in enumerate(self.train_metrics.values()):
+            log_dict[f"train_{self.metrics[i]}"] = metric_value.compute()
+        self.log_dict(log_dict,
+                    prog_bar=True)
+        
+        # print the metrics
+        str = "\n[Train] "
+        for key, value in self.trainer.logged_metrics.items():
+            if key.startswith("train"):
+                str += f"{key}: {value:.3f} "
+        print(str + '\n')
+
+        # reset the metrics
+        self.__reset_metric()
+        
+    def validation_step(self, batch: Tuple[torch.Tensor],
+                        batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        # calculate feat y_pred
+        feat = self.decoder(x)
+        y_hat = self.classifier(feat)
+
+        # get loss (pred_loss,center loss,total_loss)
+        pre_loss = self.ce_fn(y_hat, y)
+        centerloss = self.center_loss(y, feat)
+        loss = pre_loss + self.lammda * centerloss
+
+        # update  metrics
+        self.val_loss.update(loss)
+        self.center_loss_val.update(centerloss)
+        self.predict_loss_val.update(pre_loss)
+        self.val_metrics.update(y_hat, y)
+
+        # log loss
+        log_dict = {"val_loss" : self.val_loss.compute()}
+        # log metrics
+        for i, metric_value in enumerate(self.val_metrics.values()):
+            log_dict[f"val_{self.metrics[i]}"] =  metric_value.compute()
+            self.log_dict(log_dict,
+                          prog_bar=True,
+                          on_epoch=True,
+                          on_step=False,
+                          logger=True)
+        
+    
+        
+        return loss
+        
+
+    def on_validation_epoch_end(self) -> None:
+        # log loss
+        log_dict = {"val_loss" : self.val_loss.compute() }
+        # log classfication metrics
+        for i, metric_value in enumerate(self.val_metrics.values()):
+            log_dict[f"val_{self.metrics[i]}"] = metric_value.compute()
+        self.log_dict(log_dict,
+                    prog_bar=True,
+                    on_epoch=True,
+                     on_step=False,
+                     logger=True)
+        # reset the metrics
+        self.__reset_metric("val")
+        str = "\n[Val] "
+        for key, value in self.trainer.logged_metrics.items():
+            if key.startswith("val"):
+                str += f"{key}: {value:.3f} "
+        print(str + '\n')
+    
+
+
+    def test_step(self, batch: Tuple[torch.Tensor],
+                  batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        # centerloss
+        feat = self.decoder(x) 
+        centerloss = self.center_loss(y,feat)
+
+        # predict loss
+        y_hat = self.classifier(feat)
+        pre_loss = self.ce_fn(y_hat, y)
+    
+        #Total loss 
+        loss = pre_loss +self.lammda *centerloss
+
+        self.test_loss.update(loss)
+        self.center_loss_test.update(centerloss)
+        self.predict_loss_test(pre_loss)
+        self.test_metrics.update(y_hat, y)
+        return loss
+    
+    def on_test_epoch_end(self) -> None:
+        log_dict = {"test_loss" : self.test_loss.compute()}
+        # log classfication metrics
+        for i, metric_value in enumerate(self.test_metrics.values()):
+            log_dict[f"test_{self.metrics[i]}"] = metric_value.compute()
+        self.log_dict(log_dict,
+                    prog_bar=True,
+                    on_epoch=True,
+                     on_step=False,
+                     logger=True)
+        # reset the metrics
+        self.__reset_metric("test")
+        str = "\n[Test] "
+        for key, value in self.trainer.logged_metrics.items():
+            if key.startswith("test"):
+                str += f"{key}: {value:.3f} "
+        print(str + '\n')
+        
+
+        self.test_loss.reset()
+        self.test_metrics.reset()
+
+    def configure_optimizers(self):
+        parameters = list(self.decoder.parameters())
+        parameters.extend(list(self.classifier.parameters()))
+        trainable_parameters = list(
+            filter(lambda p: p.requires_grad, parameters))
+        model_optimizer = torch.optim.Adam(trainable_parameters,
+                                          lr=self.lr,
+                                          weight_decay=self.weight_decay)
+
+        center_optimizer = torch.optim.SGD(self.center_loss.parameters(), lr=0.01)
+        return center_optimizer, model_optimizer
\ No newline at end of file
diff --git a/torcheeg/trainers/imbalance/__init__.py b/torcheeg/trainers/imbalance/__init__.py
new file mode 100644
index 0000000..c71dfe6
--- /dev/null
+++ b/torcheeg/trainers/imbalance/__init__.py
@@ -0,0 +1,5 @@
+from .la import LALossTrainer
+from .ldam import LDAMLossTrainer
+from .eq import EQLossTrainer
+from .focal import FocalLossTrainer
+from .wce import WCELossTrainer
\ No newline at end of file
diff --git a/torcheeg/trainers/imbalance/eq.py b/torcheeg/trainers/imbalance/eq.py
new file mode 100644
index 0000000..a3c3ea0
--- /dev/null
+++ b/torcheeg/trainers/imbalance/eq.py
@@ -0,0 +1,152 @@
+from typing import List, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.utils.data import DataLoader
+
+from ..classifier import ClassifierTrainer
+
+
+class EQLoss(nn.Module):
+
+    def __init__(self,
+                 class_frequency: List[int],
+                 gamma: float = 0.9,
+                 lambd: float = 0.005):
+        '''
+        Equalization loss for imbalanced datasets.
+    
+        - Paper: Tan J, Wang C, Li B, et al. Equalization loss for long-tailedobject recognition[C]//Proceedings of the IEEE/CVF conference on computervision and pattern recognition. 2020: 11662-11671.
+        - URL: https://openaccess.thecvf.com/content_CVPR_2020/papersTan_Equalization_Loss_for_Long-Tailed_Object_Recognition_CVPR_2020_paper.pdf
+        - Related Project: https://github.com/tztztztztz/eql.detectron2
+
+        Args:
+            class_frequency (List[int]): The frequency of each class in the dataset.
+            gamma (float): The gamma parameter. (default: :obj:`0.9`)
+            lambd (float): The lambd parameter. (default: :obj:`0.005`)
+        '''
+        super(EQLoss, self).__init__()
+        class_frequency = torch.tensor(class_frequency)
+        self.register_buffer('class_frequency', class_frequency)
+        self.gamma = gamma
+        self.lambd = lambd
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        onehot_class = F.one_hot(target,
+                                 num_classes=self.class_frequency.shape[0])
+        beta = (torch.rand(onehot_class.size()) < self.gamma).to(input.device)
+        w = 1 - beta * self.lambd * (1 - onehot_class)
+        p = torch.exp(input) / (
+            (w * torch.exp(input)).sum(axis=1, keepdims=True))
+        loss = F.nll_loss(torch.log(p), target)
+        return loss
+
+
+class EQLossTrainer(ClassifierTrainer):
+    r'''
+        A trainer class for EEG classification with Equalization (EQ) loss for imbalanced datasets.
+    
+        - Paper: Tan J, Wang C, Li B, et al. Equalization loss for long-tailedobject recognition[C]//Proceedings of the IEEE/CVF conference on computervision and pattern recognition. 2020: 11662-11671.
+        - URL: https://openaccess.thecvf.com/content_CVPR_2020/papersTan_Equalization_Loss_for_Long-Tailed_Object_Recognition_CVPR_2020_paper.pdf
+        - Related Project: https://github.com/tztztztztz/eql.detectron2
+
+        .. code-block:: python
+
+            trainer = EQLossTrainer(model, num_classes=2, class_frequency=train_loader)
+            trainer.fit(train_loader, val_loader)
+            trainer.test(test_loader)
+
+            trainer = EQLossTrainer(model, num_classes=2, class_frequency=[10, 20], gamma=0.9, lambd=0.005)
+            trainer.fit(train_loader, val_loader)
+            trainer.test(test_loader)
+
+        Args:
+            model (nn.Module): The classification model, and the dimension of its output should be equal to the number of categories in the dataset. The output layer does not need to have a softmax activation function.
+            num_classes (int): The number of classes in the dataset.
+            class_frequency (List[int] or Dataloader): The frequency of each class in the dataset. It can be a list of integers or a dataloader to calculate the frequency of each class in the dataset, traversing the data batch (:obj:`torch.utils.data.dataloader.DataLoader`, :obj:`torch_geometric.loader.DataLoader`, etc). (default: :obj:`None`)
+            gamma (float): The gamma parameter. (default: :obj:`0.9`)
+            lambd (float): The lambd parameter. (default: :obj:`0.005`)
+            lr (float): The learning rate. (default: :obj:`0.001`)
+            weight_decay (float): The weight decay. (default: :obj:`0.0`)
+            devices (int): The number of devices to use. (default: :obj:`1`)
+            accelerator (str): The accelerator to use. Available options are: 'cpu', 'gpu'. (default: :obj:`"cpu"`)
+            metrics (list of str): The metrics to use. Available options are: 'precision', 'recall', 'f1score', 'accuracy'. (default: :obj:`["accuracy"]`)
+        
+        .. automethod:: fit
+        .. automethod:: test
+    '''
+
+    def __init__(self,
+                 model: nn.Module,
+                 num_classes: int,
+                 class_frequency: Union[List[int], DataLoader],
+                 gamma: float = 0.9,
+                 lambd: float = 0.005,
+                 lr: float = 1e-3,
+                 weight_decay: float = 0.0,
+                 devices: int = 1,
+                 accelerator: str = "cpu",
+                 metrics: List[str] = ["accuracy"]):
+
+        super().__init__(model, num_classes, lr, weight_decay, devices,
+                         accelerator, metrics)
+        self.gamma = gamma
+        self.lambd = lambd
+        self.class_frequency = class_frequency
+
+        if isinstance(class_frequency, DataLoader):
+            _class_frequency = [0] * self.num_classes
+            for _, y in class_frequency:
+                assert y < self.num_classes, f"The label in class_frequency ({y}) is out of range 0-{self.num_classes-1}."
+                _class_frequency[y] += 1
+            self._class_frequency = _class_frequency
+        else:
+            self._class_frequency = class_frequency
+
+        self.eq_fn = EQLoss(self._class_frequency, self.gamma, self.lambd)
+
+    def training_step(self, batch: Tuple[torch.Tensor],
+                      batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        y_hat = self(x)
+        loss = self.eq_fn(y_hat, y)
+
+        # log to prog_bar
+        self.log("train_loss",
+                 self.train_loss(loss),
+                 prog_bar=True,
+                 on_epoch=False,
+                 logger=False,
+                 on_step=True)
+
+        for i, metric_value in enumerate(self.train_metrics.values()):
+            self.log(f"train_{self.metrics[i]}",
+                     metric_value(y_hat, y),
+                     prog_bar=True,
+                     on_epoch=False,
+                     logger=False,
+                     on_step=True)
+
+        return loss
+
+    def validation_step(self, batch: Tuple[torch.Tensor],
+                        batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        y_hat = self(x)
+        loss = self.eq_fn(y_hat, y)
+
+        self.val_loss.update(loss)
+        self.val_metrics.update(y_hat, y)
+        return loss
+
+    def test_step(self, batch: Tuple[torch.Tensor],
+                  batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        y_hat = self(x)
+        loss = self.eq_fn(y_hat, y)
+
+        self.test_loss.update(loss)
+        self.test_metrics.update(y_hat, y)
+        return loss
\ No newline at end of file
diff --git a/torcheeg/trainers/imbalance/focal.py b/torcheeg/trainers/imbalance/focal.py
new file mode 100644
index 0000000..dccd802
--- /dev/null
+++ b/torcheeg/trainers/imbalance/focal.py
@@ -0,0 +1,178 @@
+from typing import List, Tuple, Union
+
+import numpy as np
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+
+from ..classifier import ClassifierTrainer
+
+
+class FocalLoss(nn.Module):
+
+    def __init__(self,
+                 gamma: float = 1.0,
+                 weight: Tensor = None,
+                 reduction: str = 'mean'):
+        '''
+        Focal loss for imbalanced datasets.
+
+        - Paper: Lin T Y, Goyal P, Girshick R, et al. Focal loss for dense object detection[C]//Proceedings of the IEEE international conference on computer vision. 2017: 2980-2988.
+        - URL: https://openaccess.thecvf.com/content_ICCV_2017/papers/Lin_Focal_Loss_for_ICCV_2017_paper.pdf
+        - Related Project: https://github.com/clcarwin/focal_loss_pytorch
+
+        Args:
+            gamma (float): The gamma parameter. (default: :obj:`1.0`)
+        '''
+        super(FocalLoss, self).__init__()
+        self.gamma = gamma
+        self.register_buffer('weight', weight)
+        self.reduction = reduction
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        ce_loss = F.cross_entropy(input, target, reduction='none')
+        p = torch.exp(-ce_loss)
+        loss = (1 - p)**self.gamma * ce_loss
+
+        if self.reduction == 'mean':
+            return loss.mean()
+        elif self.reduction == 'sum':
+            return loss.sum()
+        else:
+            return loss
+        
+class FocalLossTrainer(ClassifierTrainer):
+    r'''
+        A trainer class for EEG classification with Focal loss for imbalanced datasets.
+
+        - Paper: Lin T Y, Goyal P, Girshick R, et al. Focal loss for dense object detection[C]//Proceedings of the IEEE international conference on computer vision. 2017: 2980-2988.
+        - URL: https://openaccess.thecvf.com/content_ICCV_2017/papers/Lin_Focal_Loss_for_ICCV_2017_paper.pdf
+        - Related Project: https://github.com/clcarwin/focal_loss_pytorch
+
+        .. code-block:: python
+
+            trainer = FocalLossTrainer(model, num_classes=2, class_frequency=train_loader)
+            trainer.fit(train_loader, val_loader)
+            trainer.test(test_loader)
+
+            trainer = FocalLossTrainer(model, num_classes=2, class_frequency=[10, 20], gamma=1.0)
+            trainer.fit(train_loader, val_loader)
+            trainer.test(test_loader)
+
+        Args:
+            model (nn.Module): The classification model, and the dimension of its output should be equal to the number of categories in the dataset. The output layer does not need to have a softmax activation function.
+            num_classes (int): The number of classes in the dataset.
+            class_frequency (List[int] or Dataloader): The frequency of each class in the dataset. It can be a list of integers or a dataloader to calculate the frequency of each class in the dataset, traversing the data batch (:obj:`torch.utils.data.dataloader.DataLoader`, :obj:`torch_geometric.loader.DataLoader`, etc). (default: :obj:`None`)
+            gamma (float): The gamma parameter. (default: :obj:`1.0`)
+            rule (str): The rule to adjust the weight of each class. Available options are: 'none', 'reweight', 'drw' (deferred re-balancing optimization schedule). (default: :obj:`none`)
+            beta_reweight (float): The beta parameter for reweighting. It is only used when :obj:`rule` is 'reweight' or 'drw'. (default: :obj:`0.9999`)
+            drw_epochs (int): The number of epochs to use DRW. It is only used when :obj:`rule` is 'drw'. (default: :obj:`160`)
+            lr (float): The learning rate. (default: :obj:`0.001`)
+            weight_decay (float): The weight decay. (default: :obj:`0.0`)
+            devices (int): The number of devices to use. (default: :obj:`1`)
+            accelerator (str): The accelerator to use. Available options are: 'cpu', 'gpu'. (default: :obj:`"cpu"`)
+    '''
+
+    def __init__(self,
+                 model: nn.Module,
+                 num_classes: int,
+                 class_frequency: Union[List[int], DataLoader],
+                 gamma: float = 0.5,
+                 rule: str = "reweight",
+                 beta_reweight: float = 0.9999,
+                 drw_epochs: int = 160,
+                 lr: float = 1e-3,
+                 weight_decay: float = 0.0,
+                 devices: int = 1,
+                 accelerator: str = "cpu",
+                 metrics: List[str] = ["accuracy"]):
+        super().__init__(model, num_classes, lr, weight_decay, devices,
+                         accelerator, metrics)
+        self.gamma = gamma
+        self.class_frequency = class_frequency
+        self.rule = rule
+        self.beta_reweight = beta_reweight
+        self.drw_epochs = drw_epochs
+
+        if isinstance(class_frequency, DataLoader):
+            _class_frequency = [0] * self.num_classes
+            for _, y in class_frequency:
+                assert y < self.num_classes, f"The label in class_frequency ({y}) is out of range 0-{self.num_classes-1}."
+                _class_frequency[y] += 1
+            self._class_frequency = _class_frequency
+        else:
+            self._class_frequency = class_frequency
+
+        assert self.rule in ["none", "reweight",
+                             "drw"], f"Unsupported rule: {self.rule}."
+
+        if self.rule == "none":
+            _weight = None
+            self._weight = _weight
+        elif self.rule == "reweight":
+            effective_num = 1.0 - np.power(self.beta_reweight, self._class_frequency)
+            _weight = (1.0 - self.beta_reweight) / np.array(effective_num)
+            _weight = _weight / np.sum(_weight) * self.num_classes
+            self._weight = torch.tensor(_weight).float()
+        else:
+            _weight = [1.0] * self.num_classes
+            effective_num = 1.0 - np.power(self.beta_reweight, self._class_frequency)
+            _drw_weight = (1.0 - self.beta_reweight) / np.array(effective_num)
+            _drw_weight = _drw_weight / np.sum(_drw_weight) * self.num_classes
+            self._drw_weight = torch.tensor(_drw_weight).float()
+            self._weight = torch.tensor(_weight).float()
+
+        self.focal_fn = FocalLoss(gamma=self.gamma, weight=self._weight)
+
+    def on_train_epoch_start(self) -> None:
+        # get epoch
+        epoch = self.current_epoch
+        if epoch == self.drw_epochs and self.rule == "drw":
+            # reset the weight buffer in FocalLoss
+            self.focal_fn = FocalLoss(gamma=self.gamma, weight=self._drw_weight).to(self.device)
+        return super().on_train_epoch_start()
+
+    def training_step(self, batch: Tuple[torch.Tensor],
+                      batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        y_hat = self(x)
+        loss = self.focal_fn(y_hat, y)
+
+        # log to prog_bar
+        self.log("train_loss",
+                 self.train_loss(loss),
+                 prog_bar=True,
+                 on_epoch=False,
+                 logger=False,
+                 on_step=True)
+
+        for i, metric_value in enumerate(self.train_metrics.values()):
+            self.log(f"train_{self.metrics[i]}",
+                     metric_value(y_hat, y),
+                     prog_bar=True,
+                     on_epoch=False,
+                     logger=False,
+                     on_step=True)
+
+        return loss
+
+    def validation_step(self, batch: Tuple[torch.Tensor],
+                        batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        y_hat = self(x)
+        loss = self.focal_fn(y_hat, y)
+
+        self.val_loss.update(loss)
+        self.val_metrics.update(y_hat, y)
+        return loss
+
+    def test_step(self, batch: Tuple[torch.Tensor],
+                  batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        y_hat = self(x)
+        loss = self.focal_fn(y_hat, y)
+
+        self.test_loss.update(loss)
+        self.test_metrics.update(y_hat, y)
+        return loss
\ No newline at end of file
diff --git a/torcheeg/trainers/imbalance/la.py b/torcheeg/trainers/imbalance/la.py
new file mode 100644
index 0000000..c646a5d
--- /dev/null
+++ b/torcheeg/trainers/imbalance/la.py
@@ -0,0 +1,146 @@
+from typing import List, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.utils.data import DataLoader
+
+from ..classifier import ClassifierTrainer
+
+
+class LALoss(nn.Module):
+
+    def __init__(self, class_frequency: List[int], tau=1.0, eps=1e-12):
+        '''
+        Logit-adjusted (LA) loss for imbalanced datasets.
+
+        - Paper: Menon A K, Jayasumana S, Rawat A S, et al. Long-tail learning via logit adjustment[J]. arXiv preprint arXiv:2007.07314, 2020.
+        - URL: https://arxiv.org/abs/2007.07314
+        - Related Project: https://github.com/Chumsy0725/logit-adj-pytorch
+
+        Args:
+            class_frequency (List[int]): The frequency of each class in the dataset.
+            tau (float): The temperature parameter. (default: :obj:`1.0`)
+            eps (float): The epsilon parameter. (default: :obj:`1e-12`)
+        '''
+        super(LALoss, self).__init__()
+        class_frequency = torch.tensor(class_frequency)
+        self.register_buffer('class_frequency', class_frequency)
+
+        label_probability = class_frequency / class_frequency.sum()
+        adjustments = tau * torch.log(label_probability + eps)
+        adjustments = adjustments.reshape(1, -1)
+        self.register_buffer('adjustments', adjustments.float())
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        input += self.adjustments
+        return F.cross_entropy(input, target)
+
+
+class LALossTrainer(ClassifierTrainer):
+    r'''
+        A trainer class for EEG classification with Logit-adjusted (LA) loss for imbalanced datasets.
+
+        - Paper: Menon A K, Jayasumana S, Rawat A S, et al. Long-tail learning via logit adjustment[J]. arXiv preprint arXiv:2007.07314, 2020.
+        - URL: https://arxiv.org/abs/2007.07314
+        - Related Project: https://github.com/Chumsy0725/logit-adj-pytorch
+
+        .. code-block:: python
+
+            trainer = LALossTrainer(model, num_classes=2, class_frequency=train_loader)
+            trainer.fit(train_loader, val_loader)
+            trainer.test(test_loader)
+
+            trainer = LALossTrainer(model, num_classes=2, class_frequency=[10, 20], tau=1.0)
+            trainer.fit(train_loader, val_loader)
+            trainer.test(test_loader)
+
+        Args:
+            model (nn.Module): The classification model, and the dimension of its output should be equal to the number of categories in the dataset. The output layer does not need to have a softmax activation function.
+            num_classes (int): The number of classes in the dataset.
+            class_frequency (List[int] or Dataloader): The frequency of each class in the dataset. It can be a list of integers or a dataloader to calculate the frequency of each class in the dataset, traversing the data batch (:obj:`torch.utils.data.dataloader.DataLoader`, :obj:`torch_geometric.loader.DataLoader`, etc). (default: :obj:`None`)
+            tau (float): The temperature parameter. (default: :obj:`1.0`)
+            eps (float): The epsilon parameter. (default: :obj:`1e-12`)
+            lr (float): The learning rate. (default: :obj:`0.001`)
+            weight_decay (float): The weight decay. (default: :obj:`0.0`)
+            devices (int): The number of devices to use. (default: :obj:`1`)
+            accelerator (str): The accelerator to use. Available options are: 'cpu', 'gpu'. (default: :obj:`"cpu"`)
+            metrics (list of str): The metrics to use. Available options are: 'precision', 'recall', 'f1score', 'accuracy'. (default: :obj:`["accuracy"]`)
+        
+        .. automethod:: fit
+        .. automethod:: test
+    '''
+
+    def __init__(self,
+                 model: nn.Module,
+                 num_classes: int,
+                 class_frequency: Union[List[int], DataLoader],
+                 tau: float = 1.0,
+                 eps: float = 1e-12,
+                 lr: float = 1e-3,
+                 weight_decay: float = 0.0,
+                 devices: int = 1,
+                 accelerator: str = "cpu",
+                 metrics: List[str] = ["accuracy"]):
+
+        super().__init__(model, num_classes, lr, weight_decay, devices,
+                         accelerator, metrics)
+        self.tau = tau
+        self.eps = eps
+        self.class_frequency = class_frequency
+
+        if isinstance(class_frequency, DataLoader):
+            _class_frequency = [0] * self.num_classes
+            for _, y in class_frequency:
+                assert y < self.num_classes, f"The label in class_frequency ({y}) is out of range 0-{self.num_classes-1}."
+                _class_frequency[y] += 1
+            self._class_frequency = _class_frequency
+        else:
+            self._class_frequency = class_frequency
+
+        self.la_fn = LALoss(self._class_frequency, self.tau, self.eps)
+
+    def training_step(self, batch: Tuple[torch.Tensor],
+                      batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        y_hat = self(x)
+        loss = self.la_fn(y_hat, y)
+
+        # log to prog_bar
+        self.log("train_loss",
+                 self.train_loss(loss),
+                 prog_bar=True,
+                 on_epoch=False,
+                 logger=False,
+                 on_step=True)
+
+        for i, metric_value in enumerate(self.train_metrics.values()):
+            self.log(f"train_{self.metrics[i]}",
+                     metric_value(y_hat, y),
+                     prog_bar=True,
+                     on_epoch=False,
+                     logger=False,
+                     on_step=True)
+
+        return loss
+
+    def validation_step(self, batch: Tuple[torch.Tensor],
+                        batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        y_hat = self(x)
+        loss = self.la_fn(y_hat, y)
+
+        self.val_loss.update(loss)
+        self.val_metrics.update(y_hat, y)
+        return loss
+
+    def test_step(self, batch: Tuple[torch.Tensor],
+                  batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        y_hat = self(x)
+        loss = self.la_fn(y_hat, y)
+
+        self.test_loss.update(loss)
+        self.test_metrics.update(y_hat, y)
+        return loss
\ No newline at end of file
diff --git a/torcheeg/trainers/imbalance/ldam.py b/torcheeg/trainers/imbalance/ldam.py
new file mode 100644
index 0000000..c335f20
--- /dev/null
+++ b/torcheeg/trainers/imbalance/ldam.py
@@ -0,0 +1,201 @@
+from typing import Any, List, Tuple, Union
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+
+from ..classifier import ClassifierTrainer
+
+
+class LDAMLoss(nn.Module):
+
+    def __init__(self,
+                 class_frequency: List[int],
+                 max_margin: float = 0.5,
+                 weight: Tensor = None,
+                 scaling: float = 30):
+        '''
+        Label-distribution-aware margin (LDAM) loss for imbalanced datasets.
+
+        - Paper: Cao K, Wei C, Gaidon A, et al. Learning imbalanced datasets with label-distribution-aware margin loss[J]. Advances in neural information processing systems, 2019, 32.
+        - URL: https://proceedings.neurips.cc/paper_files/paper/2019/file/621461af90cadfdaf0e8d4cc25129f91-Paper.pdf
+        - Related Project: https://github.com/kaidic/LDAM-DRW
+
+        Args:
+            class_frequency (List[int]): The frequency of each class in the dataset.
+            max_margin (float): The maximum margin. (default: :obj:`0.5`)
+            weight (Tensor): The weight of each class. (default: :obj:`None`)
+            scaling (float): The scaling factor. (default: :obj:`30`)
+        '''
+        super(LDAMLoss, self).__init__()
+        margin_list = 1.0 / np.sqrt(np.sqrt(class_frequency))
+        margin_list = margin_list * (max_margin / np.max(margin_list))
+        self.register_buffer('margin_list', torch.tensor(margin_list).float())
+        assert scaling > 0, "scaling should be greater than 0."
+        self.scaling = scaling
+        if not weight is None:
+            self.register_buffer('weight', weight)
+        else:
+            self.weight = None
+
+    def forward(self, input: Tensor, target: Tensor) -> Tensor:
+        index = torch.zeros_like(input)
+        index.scatter_(1, target.data.view(-1, 1), 1)
+
+        index_float = index.float()
+        index_bool = index.bool()
+
+        batch_m = torch.matmul(self.margin_list[None, :],
+                               index_float.transpose(0, 1))
+        batch_m = batch_m.view((-1, 1))
+        x_m = input - batch_m
+
+        output = torch.where(index_bool, x_m, input)
+        return F.cross_entropy(self.scaling * output,
+                               target,
+                               weight=self.weight)
+
+
+class LDAMLossTrainer(ClassifierTrainer):
+    r'''
+        A trainer class for EEG classification with Label-distribution-aware margin (LDAM) loss for imbalanced datasets.
+
+        - Paper: Cao K, Wei C, Gaidon A, et al. Learning imbalanced datasets with label-distribution-aware margin loss[J]. Advances in neural information processing systems, 2019, 32.
+        - URL: https://proceedings.neurips.cc/paper_files/paper/2019/file/621461af90cadfdaf0e8d4cc25129f91-Paper.pdf
+        - Related Project: https://github.com/kaidic/LDAM-DRW
+
+        .. code-block:: python
+
+            trainer = LDAMLossTrainer(model, num_classes=2, class_frequency=train_loader)
+            trainer.fit(train_loader, val_loader)
+            trainer.test(test_loader)
+
+            trainer = LDAMLossTrainer(model, num_classes=2, class_frequency=[10, 20], max_margin=0.5, scaling=30)
+            trainer.fit(train_loader, val_loader)
+            trainer.test(test_loader)
+
+        Args:
+            model (nn.Module): The classification model, and the dimension of its output should be equal to the number of categories in the dataset. The output layer does not need to have a softmax activation function.
+            num_classes (int): The number of classes in the dataset.
+            class_frequency (List[int] or Dataloader): The frequency of each class in the dataset. It can be a list of integers or a dataloader to calculate the frequency of each class in the dataset, traversing the data batch (:obj:`torch.utils.data.dataloader.DataLoader`, :obj:`torch_geometric.loader.DataLoader`, etc). (default: :obj:`None`)
+            max_margin (float): The maximum margin. (default: :obj:`0.5`)
+            rule (str): The rule to adjust the weight of each class. Available options are: 'none', 'reweight', 'drw' (deferred re-balancing optimization schedule). (default: :obj:`'none'`)
+            beta_reweight (float): The beta parameter for reweighting. It is only used when :obj:`rule` is 'reweight' or 'drw'. (default: :obj:`0.9999`)
+            drw_epochs (int): The number of epochs to use DRW. It is only used when :obj:`rule` is 'drw'. (default: :obj:`160`)
+            scaling (float): The scaling factor. (default: :obj:`30`)
+            lr (float): The learning rate. (default: :obj:`0.001`)
+            weight_decay (float): The weight decay. (default: :obj:`0.0`)
+            devices (int): The number of devices to use. (default: :obj:`1`)
+            accelerator (str): The accelerator to use. Available options are: 'cpu', 'gpu'. (default: :obj:`"cpu"`)
+    '''
+
+    def __init__(self,
+                 model: nn.Module,
+                 num_classes: int,
+                 class_frequency: Union[List[int], DataLoader],
+                 max_margin: float = 0.5,
+                 scaling: float = 30,
+                 rule: str = "none",
+                 beta_reweight: float = 0.9999,
+                 drw_epochs: int = 160,
+                 lr: float = 1e-3,
+                 weight_decay: float = 0.0,
+                 devices: int = 1,
+                 accelerator: str = "cpu",
+                 metrics: List[str] = ["accuracy"]):
+        super().__init__(model, num_classes, lr, weight_decay, devices,
+                         accelerator, metrics)
+        self.max_margin = max_margin
+        self.scaling = scaling
+        self.class_frequency = class_frequency
+        self.rule = rule
+        self.beta_reweight = beta_reweight
+        self.drw_epochs = drw_epochs
+
+        if isinstance(class_frequency, DataLoader):
+            _class_frequency = [0] * self.num_classes
+            for _, y in class_frequency:
+                assert y < self.num_classes, f"The label in class_frequency ({y}) is out of range 0-{self.num_classes-1}."
+                _class_frequency[y] += 1
+            self._class_frequency = _class_frequency
+        else:
+            self._class_frequency = class_frequency
+
+        assert self.rule in ["none", "reweight",
+                             "drw"], f"Unsupported rule: {self.rule}."
+
+        if self.rule == 'none':
+            _weight = None
+            self._weight = None
+        elif self.rule == "reweight":
+            effective_num = 1.0 - np.power(self.beta_reweight, self._class_frequency)
+            _weight = (1.0 - self.beta_reweight) / np.array(effective_num)
+            _weight = _weight / np.sum(_weight) * self.num_classes
+            self._weight = torch.tensor(_weight).float()
+        else:
+            _weight = [1.0] * self.num_classes
+            effective_num = 1.0 - np.power(self.beta_reweight, self._class_frequency)
+            _drw_weight = (1.0 - self.beta_reweight) / np.array(effective_num)
+            _drw_weight = _drw_weight / np.sum(_drw_weight) * self.num_classes
+            self._drw_weight = torch.tensor(_drw_weight).float()
+            self._weight = torch.tensor(_weight).float()
+
+        self.ldam_fn = LDAMLoss(self._class_frequency, max_margin, self._weight,
+                                scaling)
+
+    def on_train_epoch_start(self) -> None:
+        # get epoch
+        epoch = self.current_epoch
+        if epoch == self.drw_epochs and self.rule == "drw":
+            # reset the weight buffer in LDAMLoss
+            self.ldam_fn = LDAMLoss(self._class_frequency, self.max_margin,
+                                    self._drw_weight,
+                                    self.scaling).to(self.device)
+        return super().on_train_epoch_start()
+
+    def training_step(self, batch: Tuple[torch.Tensor],
+                      batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        y_hat = self(x)
+        loss = self.ldam_fn(y_hat, y)
+
+        # log to prog_bar
+        self.log("train_loss",
+                 self.train_loss(loss),
+                 prog_bar=True,
+                 on_epoch=False,
+                 logger=False,
+                 on_step=True)
+
+        for i, metric_value in enumerate(self.train_metrics.values()):
+            self.log(f"train_{self.metrics[i]}",
+                     metric_value(y_hat, y),
+                     prog_bar=True,
+                     on_epoch=False,
+                     logger=False,
+                     on_step=True)
+
+        return loss
+
+    def validation_step(self, batch: Tuple[torch.Tensor],
+                        batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        y_hat = self(x)
+        loss = self.ldam_fn(y_hat, y)
+
+        self.val_loss.update(loss)
+        self.val_metrics.update(y_hat, y)
+        return loss
+
+    def test_step(self, batch: Tuple[torch.Tensor],
+                  batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        y_hat = self(x)
+        loss = self.ldam_fn(y_hat, y)
+
+        self.test_loss.update(loss)
+        self.test_metrics.update(y_hat, y)
+        return loss
\ No newline at end of file
diff --git a/torcheeg/trainers/imbalance/wce.py b/torcheeg/trainers/imbalance/wce.py
new file mode 100644
index 0000000..a25fff5
--- /dev/null
+++ b/torcheeg/trainers/imbalance/wce.py
@@ -0,0 +1,133 @@
+from typing import List, Tuple, Union
+
+import numpy as np
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+
+from ..classifier import ClassifierTrainer
+        
+class WCELossTrainer(ClassifierTrainer):
+    r'''
+        A trainer class for EEG classification with Weighted Cross Entropy (WCE) loss for imbalanced datasets.
+
+        .. code-block:: python
+
+            trainer = WCELossTrainer(model, num_classes=2, class_frequency=train_loader)
+            trainer.fit(train_loader, val_loader)
+            trainer.test(test_loader)
+
+            trainer = WCELossTrainer(model, num_classes=2, class_frequency=[10, 20])
+            trainer.fit(train_loader, val_loader)
+            trainer.test(test_loader)
+
+        Args:
+            model (nn.Module): The classification model, and the dimension of its output should be equal to the number of categories in the dataset. The output layer does not need to have a softmax activation function.
+            num_classes (int): The number of classes in the dataset.
+            class_frequency (List[int] or Dataloader): The frequency of each class in the dataset. It can be a list of integers or a dataloader to calculate the frequency of each class in the dataset, traversing the data batch (:obj:`torch.utils.data.dataloader.DataLoader`, :obj:`torch_geometric.loader.DataLoader`, etc). (default: :obj:`None`)
+            gamma (float): The gamma parameter. (default: :obj:`1.0`)
+            rule (str): The rule to adjust the weight of each class. Available options are: 'reweight', 'drw' (deferred re-balancing optimization schedule). (default: :obj:`reweight`)
+            beta_reweight (float): The beta parameter for reweighting. It is only used when :obj:`rule` is 'reweight' or 'drw'. (default: :obj:`0.9999`)
+            drw_epochs (int): The number of epochs to use DRW. It is only used when :obj:`rule` is 'drw'. (default: :obj:`160`)
+            lr (float): The learning rate. (default: :obj:`0.001`)
+            weight_decay (float): The weight decay. (default: :obj:`0.0`)
+            devices (int): The number of devices to use. (default: :obj:`1`)
+            accelerator (str): The accelerator to use. Available options are: 'cpu', 'gpu'. (default: :obj:`"cpu"`)
+    '''
+
+    def __init__(self,
+                 model: nn.Module,
+                 num_classes: int,
+                 class_frequency: Union[List[int], DataLoader],
+                 rule: str = "reweight",
+                 beta_reweight: float = 0.9999,
+                 drw_epochs: int = 160,
+                 lr: float = 1e-3,
+                 weight_decay: float = 0.0,
+                 devices: int = 1,
+                 accelerator: str = "cpu",
+                 metrics: List[str] = ["accuracy"]):
+        super().__init__(model, num_classes, lr, weight_decay, devices,
+                         accelerator, metrics)
+        self.class_frequency = class_frequency
+        self.rule = rule
+        self.beta_reweight = beta_reweight
+        self.drw_epochs = drw_epochs
+
+        if isinstance(class_frequency, DataLoader):
+            _class_frequency = [0] * self.num_classes
+            for _, y in class_frequency:
+                assert y < self.num_classes, f"The label in class_frequency ({y}) is out of range 0-{self.num_classes-1}."
+                _class_frequency[y] += 1
+            self._class_frequency = _class_frequency
+        else:
+            self._class_frequency = class_frequency
+
+        assert self.rule in ["none", "reweight",
+                             "drw"], f"Unsupported rule: {self.rule}."
+
+        if self.rule == "reweight":
+            effective_num = 1.0 - np.power(self.beta_reweight, self._class_frequency)
+            _weight = (1.0 - self.beta_reweight) / np.array(effective_num)
+            _weight = _weight / np.sum(_weight) * self.num_classes
+        else:
+            _weight = [1.0] * self.num_classes
+            effective_num = 1.0 - np.power(self.beta_reweight, self._class_frequency)
+            _drw_weight = (1.0 - self.beta_reweight) / np.array(effective_num)
+            _drw_weight = _drw_weight / np.sum(_drw_weight) * self.num_classes
+            self._drw_weight = torch.tensor(_drw_weight).float()
+        self._weight = torch.tensor(_weight).float()
+
+        self.wce_fn = nn.CrossEntropyLoss(weight=self._weight)
+
+    def on_train_epoch_start(self) -> None:
+        # get epoch
+        epoch = self.current_epoch
+        if epoch == self.drw_epochs and self.rule == "drw":
+            # reset the weight buffer in CrossEntropyLoss
+            self.wce_fn = nn.CrossEntropyLoss(weight=self._drw_weight).to(self.device)
+        return super().on_train_epoch_start()
+
+    def training_step(self, batch: Tuple[torch.Tensor],
+                      batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        y_hat = self(x)
+        loss = self.wce_fn(y_hat, y)
+
+        # log to prog_bar
+        self.log("train_loss",
+                 self.train_loss(loss),
+                 prog_bar=True,
+                 on_epoch=False,
+                 logger=False,
+                 on_step=True)
+
+        for i, metric_value in enumerate(self.train_metrics.values()):
+            self.log(f"train_{self.metrics[i]}",
+                     metric_value(y_hat, y),
+                     prog_bar=True,
+                     on_epoch=False,
+                     logger=False,
+                     on_step=True)
+
+        return loss
+
+    def validation_step(self, batch: Tuple[torch.Tensor],
+                        batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        y_hat = self(x)
+        loss = self.wce_fn(y_hat, y)
+
+        self.val_loss.update(loss)
+        self.val_metrics.update(y_hat, y)
+        return loss
+
+    def test_step(self, batch: Tuple[torch.Tensor],
+                  batch_idx: int) -> torch.Tensor:
+        x, y = batch
+        y_hat = self(x)
+        loss = self.wce_fn(y_hat, y)
+
+        self.test_loss.update(loss)
+        self.test_metrics.update(y_hat, y)
+        return loss
\ No newline at end of file
diff --git a/torcheeg/transforms/numpy/band.py b/torcheeg/transforms/numpy/band.py
index e7eddd7..6a2572a 100644
--- a/torcheeg/transforms/numpy/band.py
+++ b/torcheeg/transforms/numpy/band.py
@@ -148,7 +148,7 @@ def __call__(self,
         return super().__call__(*args, eeg=eeg, baseline=baseline, **kwargs)
 
     def opt(self, eeg: np.ndarray, **kwargs) -> np.ndarray:
-        return 1 / 2 * np.log2(2 * np.pi * np.e * np.std(eeg))
+        return 1 / 2 * np.log2(2 * np.pi * np.e * np.var(eeg))
 
 
 class BandDifferentialEntropyV1(EEGTransform):


Test patch:
diff --git a/test/datasets/test_motor_imagery.py b/test/datasets/test_motor_imagery.py
new file mode 100644
index 0000000..b76e13d
--- /dev/null
+++ b/test/datasets/test_motor_imagery.py
@@ -0,0 +1,64 @@
+import os
+import random
+import shutil
+import unittest
+
+from torcheeg import transforms
+from torcheeg.datasets import BCICIV2aDataset
+
+
+class TestMotorImagery(unittest.TestCase):
+    def setUp(self):
+        if os.path.exists('./tmp_out/'):
+            shutil.rmtree('./tmp_out/')
+        os.mkdir('./tmp_out/')
+
+    def test_bciciv_2a_dataset(self):
+        io_path = f'./tmp_out/bciciv_2a_{"".join(random.sample("zyxwvutsrqponmlkjihgfedcba", 20))}'
+        root_path = './tmp_in/BCICIV_2a_mat'
+
+        dataset = BCICIV2aDataset(
+            io_path=io_path,
+            root_path=root_path,
+            online_transform=transforms.Compose([
+                                  transforms.To2d(),
+                                  transforms.ToTensor()
+                              ]),
+                              label_transform=transforms.Compose([
+                                  transforms.Select('label'),
+                                  transforms.Lambda(lambda x: x - 1)
+                              ]),
+            num_worker=4)
+
+        self.assertEqual(len(dataset), 5184)
+        first_item = dataset[0]
+        self.assertEqual(first_item[0].shape, (1, 22, 1750))
+        last_item = dataset[5183]
+        self.assertEqual(last_item[0].shape, (1, 22, 1750))
+
+        io_path = f'./tmp_out/bciciv_2a_{"".join(random.sample("zyxwvutsrqponmlkjihgfedcba", 20))}'
+        root_path = './tmp_in/BCICIV_2a_mat'
+
+        dataset = BCICIV2aDataset(
+            io_path=io_path,
+            root_path=root_path,
+            offset=1.5 * 250,
+            chunk_size=4.5 * 250,
+            online_transform=transforms.Compose([
+                                  transforms.To2d(),
+                                  transforms.ToTensor()
+                              ]),
+                              label_transform=transforms.Compose([
+                                  transforms.Select('label'),
+                                  transforms.Lambda(lambda x: x - 1)
+                              ]),
+            num_worker=4)
+
+        self.assertEqual(len(dataset), 5184)
+        first_item = dataset[0]
+        self.assertEqual(first_item[0].shape, (1, 22, 1125))
+        last_item = dataset[5183]
+        self.assertEqual(last_item[0].shape, (1, 22, 1125))
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/test/models/test_cnn.py b/test/models/test_cnn.py
index 6151a80..36e82f4 100644
--- a/test/models/test_cnn.py
+++ b/test/models/test_cnn.py
@@ -1,7 +1,8 @@
 import unittest
 
 import torch
-from torcheeg.models import (CCNN, FBCCNN, MTCNN, EEGNet, FBCNet, STNet, TSCeption, SSTEmotionNet)
+from torcheeg.models import (CCNN, FBCCNN, MTCNN, EEGNet, FBCNet, STNet,
+                             TSCeption, SSTEmotionNet, FBMSNet)
 
 
 class TestCNN(unittest.TestCase):
@@ -44,7 +45,10 @@ def test_eegnet(self):
 
     def test_stnet(self):
         eeg = torch.randn(1, 128, 9, 9)
-        model = STNet(num_classes=2, chunk_size=128, grid_size=(9, 9), dropout=0.2)
+        model = STNet(num_classes=2,
+                      chunk_size=128,
+                      grid_size=(9, 9),
+                      dropout=0.2)
         pred = model(eeg)
         self.assertEqual(tuple(pred.shape), (1, 2))
 
@@ -55,7 +59,10 @@ def test_stnet(self):
 
     def test_mtcnn(self):
         eeg = torch.randn(1, 8, 8, 9)
-        model = MTCNN(num_classes=2, in_channels=8, grid_size=(8, 9), dropout=0.2)
+        model = MTCNN(num_classes=2,
+                      in_channels=8,
+                      grid_size=(8, 9),
+                      dropout=0.2)
         pred = model(eeg)
         self.assertEqual(tuple(pred[0].shape), (1, 2))
         self.assertEqual(tuple(pred[1].shape), (1, 2))
@@ -79,7 +86,11 @@ def test_fbccnn(self):
 
     def test_fbcnet(self):
         eeg = torch.randn(1, 4, 32, 512)
-        model = FBCNet(num_classes=2, num_electrodes=32, chunk_size=512, in_channels=4, num_S=32)
+        model = FBCNet(num_classes=2,
+                       num_electrodes=32,
+                       chunk_size=512,
+                       in_channels=4,
+                       num_S=32)
         pred = model(eeg)
         self.assertEqual(tuple(pred.shape), (1, 2))
 
@@ -101,7 +112,10 @@ def test_ccnn(self):
 
     def test_sst_emotion_net(self):
         eeg = torch.randn(2, 32 + 4, 16, 16)
-        model = SSTEmotionNet(temporal_in_channels=32, spectral_in_channels=4, grid_size=(16, 16), num_classes=2)
+        model = SSTEmotionNet(temporal_in_channels=32,
+                              spectral_in_channels=4,
+                              grid_size=(16, 16),
+                              num_classes=2)
         pred = model(eeg)
         self.assertEqual(tuple(pred.shape), (2, 2))
 
@@ -110,6 +124,53 @@ def test_sst_emotion_net(self):
         pred = model(eeg)
         self.assertEqual(tuple(pred.shape), (2, 2))
 
+    def test_fbmsnet(self):
+        eeg = torch.randn(2, 9, 22, 512)
+        model = FBMSNet(in_channels=9,
+                        num_electrodes=22,
+                        chunk_size=512,
+                        num_classes=4)
+
+        code = model.decoder(eeg)
+        pred = model.classifier(code)
+  
+        self.assertEqual(tuple(pred.shape), (2, 4))
+        self.assertEqual(tuple(code.shape), (2, 1152))
+
+        eeg = eeg.cuda()
+        model = model.cuda()
+        code = model.decoder(eeg)
+        pred = model.classifier(code)
+
+
+        self.assertEqual(tuple(pred.shape), (2, 4))
+        self.assertEqual(tuple(code.shape), (2, 1152))
+
+        # other shape
+        eeg = torch.randn(2, 12, 32, 256)
+        model = FBMSNet(in_channels=12,
+                        num_electrodes=32,
+                        chunk_size=256,
+                        num_classes=3)
+        code = model.decoder(eeg)
+        pred = model.classifier(code)
+
+
+        self.assertEqual(tuple(pred.shape), (2, 3))
+        self.assertEqual(tuple(code.shape), (2, 1152))
+
+        eeg = torch.randn(2, 15, 18, 128)
+        model = FBMSNet(in_channels=15,
+                        num_electrodes=18,
+                        chunk_size=128,
+                        num_classes=5)
+        code = model.decoder(eeg)
+        pred = model.classifier(code)
+
+ 
+        self.assertEqual(tuple(pred.shape), (2, 5))
+        self.assertEqual(tuple(code.shape), (2, 1152))
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/models/test_gan.py b/test/models/test_gan.py
index 13e68f0..64117c8 100644
--- a/test/models/test_gan.py
+++ b/test/models/test_gan.py
@@ -1,8 +1,7 @@
 import unittest
 
 import torch
-from torcheeg.models import BCGenerator, BCDiscriminator, BGenerator, BDiscriminator
-
+from torcheeg.models import BCGenerator, BCDiscriminator, BGenerator, BDiscriminator,EEGfuseNet,EFDiscriminator
 
 class TestGAN(unittest.TestCase):
     def test_bgan(self):
@@ -39,6 +38,58 @@ def test_bcgan(self):
         disc_X = d_model(fake_X, y)
         self.assertEqual(tuple(disc_X.shape), (1, 1))
 
+    def test_eegfusenet(self):
+
+        g_model = EEGfuseNet(num_electrodes=20,hid_channels_gru=16,num_layers_gru=1,hid_channels_cnn=1,chunk_size=128)
+        d_model = EFDiscriminator(num_electrodes=20,hid_channels_cnn=1,chunk_size=128)
+        X = torch.rand(2,1,20,128)
+        fake_X, deep_feature = g_model(X)
+        p_real,p_fake = d_model(X),d_model(fake_X)
+
+        self.assertEqual(tuple(fake_X.shape), tuple(X.shape))
+        self.assertEqual(tuple(deep_feature.shape), (2,64))
+        self.assertEqual(tuple(p_real.shape),(2,1))
+        self.assertEqual(tuple(p_fake.shape),(2,1))
+        
+        # to gpu
+        g_model = EEGfuseNet(num_electrodes=20,hid_channels_gru=16,num_layers_gru=1,hid_channels_cnn=1,chunk_size=128).cuda(0)
+        d_model = EFDiscriminator(num_electrodes=20,hid_channels_cnn=1,chunk_size=128).cuda(0)
+        X = torch.rand(2,1,20,128).cuda(0)
+        fake_X, deep_feature = g_model(X)
+        p_real,p_fake = d_model(X),d_model(fake_X)
+
+        self.assertEqual(tuple(fake_X.shape), tuple(X.shape))
+        self.assertEqual(tuple(deep_feature.shape), (2,64))
+        self.assertEqual(tuple(p_real.shape),(2,1))
+        self.assertEqual(tuple(p_fake.shape),(2,1))
+        
+
+        # other shapes 1 
+        g_model = EEGfuseNet(num_electrodes=64,hid_channels_gru=16,num_layers_gru=1,hid_channels_cnn=1,chunk_size=512)
+        d_model = EFDiscriminator(num_electrodes=64,hid_channels_cnn=1,chunk_size=512)
+        X = torch.rand(2,1,64,512)
+        fake_X, deep_feature = g_model(X)
+        p_real,p_fake = d_model(X),d_model(fake_X)
+
+        self.assertEqual(tuple(fake_X.shape), tuple(X.shape))
+        self.assertEqual(tuple(deep_feature.shape), (2,256))
+        self.assertEqual(tuple(p_real.shape),(2,1))
+        self.assertEqual(tuple(p_fake.shape),(2,1))
+
+        
+
+        # other shapes 2
+        g_model = EEGfuseNet(num_electrodes=32,hid_channels_gru=16,num_layers_gru=1,hid_channels_cnn=1,chunk_size=384)
+        d_model = EFDiscriminator(num_electrodes=32,hid_channels_cnn=1,chunk_size=384)
+        X = torch.randn(2,1,32,384)
+        fake_X, deep_feature = g_model(X)
+        p_real,p_fake = d_model(X),d_model(fake_X)
+
+        self.assertEqual(tuple(fake_X.shape), tuple(X.shape))
+        self.assertEqual(tuple(deep_feature.shape), (2,192))
+        self.assertEqual(tuple(p_real.shape),(2,1))
+        self.assertEqual(tuple(p_fake.shape),(2,1))
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/models/test_transformer.py b/test/models/test_transformer.py
index c403ed5..ba48e8d 100644
--- a/test/models/test_transformer.py
+++ b/test/models/test_transformer.py
@@ -2,7 +2,7 @@
 
 import torch
 
-from torcheeg.models import SimpleViT, ArjunViT, VanillaTransformer, ViT, Conformer
+from torcheeg.models import SimpleViT, ArjunViT, VanillaTransformer, ViT, Conformer,ATCNet
 
 
 class TestTransformer(unittest.TestCase):
@@ -69,6 +69,30 @@ def test_conformer(self):
         model = model.cuda()
         pred = model(eeg).cpu()
         self.assertEqual(tuple(pred.shape), (2, 2))
+    
+    def test_atcnet(self):
+        eeg = torch.randn(2,1,32,1000)
+        model = ATCNet(num_classes=3,num_electrodes=32,chunk_size=1000)
+        pred = model(eeg)
+        self.assertEqual(tuple(pred.shape),(2,3))
+
+        # to gpu
+        eeg = eeg.cuda(0)
+        model = model.cuda(0)
+        pred = model(eeg)
+        self.assertEqual(tuple(pred.shape),(2,3))
+
+        # other shape
+        eeg = torch.randn(2,1,64,512)
+        model =  ATCNet(num_classes=4,num_electrodes=64,chunk_size=512)
+        pred = model(eeg)
+        self.assertEqual(tuple(pred.shape),(2,4))
+
+        # other shape
+        eeg = torch.rand(2,1,96,256)
+        model =  ATCNet(num_classes=4,num_electrodes=96,chunk_size=256)
+        pred = model(eeg)
+        self.assertEqual(tuple(pred.shape),(2,4))
 
 
 if __name__ == '__main__':
diff --git a/test/trainers/test_classifier_trainer.py b/test/trainers/test_classifier_trainer.py
index e56b6bd..f7bc739 100644
--- a/test/trainers/test_classifier_trainer.py
+++ b/test/trainers/test_classifier_trainer.py
@@ -17,7 +17,7 @@ def __len__(self):
         return self.length
 
     def __getitem__(self, i) -> int:
-        return torch.randn(120), random.randint(0, 1)
+        return torch.randn(120), random.choice([0, 1])
 
 
 class DummyModel(nn.Module):
@@ -32,9 +32,9 @@ def forward(self, x):
         return self.fc(x)
 
 
-class TestClassificationTrainer(unittest.TestCase):
+class TestClassifierTrainer(unittest.TestCase):
 
-    def test_classification_trainer(self):
+    def test_classifier_trainer(self):
         train_dataset = DummyDataset()
         val_dataset = DummyDataset()
         test_dataset = DummyDataset()
diff --git a/test/trainers/test_domain_adaption_trainer.py b/test/trainers/test_domain_adaption_trainer.py
index 1e4d995..b529d2a 100644
--- a/test/trainers/test_domain_adaption_trainer.py
+++ b/test/trainers/test_domain_adaption_trainer.py
@@ -5,10 +5,11 @@
 import torch.nn as nn
 from torch.utils.data import Dataset
 from torch.utils.data.dataloader import DataLoader
-from torcheeg.trainers import ADATrainer, CORALTrainer, DANTrainer, DANNTrainer, DDCTrainer, JANTrainer
+from torcheeg.trainers import ADATrainer, CORALTrainer, DANTrainer, DANNTrainer, DDCTrainer, JANTrainer, CenterLossTrainer
 
 
 class DummyDataset(Dataset):
+
     def __init__(self, length: int = 101):
         self.length = length
 
@@ -20,6 +21,7 @@ def __getitem__(self, i) -> int:
 
 
 class DummyModel(nn.Module):
+
     def __init__(self, in_channels=120, out_channels=2):
         super().__init__()
         self.in_channels = in_channels
@@ -31,6 +33,7 @@ def forward(self, x):
 
 
 class TestDomainAdaptionTrainer(unittest.TestCase):
+
     def test_ada_trainer(self):
         source_dataset = DummyDataset()
         target_dataset = DummyDataset()
@@ -143,6 +146,48 @@ def test_jan_trainer(self):
         trainer.fit(source_loader, target_loader, val_loader, max_epochs=1)
         trainer.test(test_loader)
 
+    def test_center_loss_trainer(self):
+
+        train_dataset = DummyDataset(length=10)
+        val_dataset = DummyDataset(length=10)
+        test_dataset = DummyDataset(length=10)
+
+        train_loader = DataLoader(train_dataset, batch_size=5)
+        val_loader = DataLoader(val_dataset, batch_size=5)
+        test_loader = DataLoader(test_dataset, batch_size=5)
+
+        decoder = DummyModel(120, 64)
+        classifier = DummyModel(64, 2)
+
+        trainer = CenterLossTrainer(decoder=decoder,
+                                    classifier=classifier,
+                                    feature_dim=64,
+                                    num_classes=2)
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+        # should catch value error for metrics 'unexpected'
+        with self.assertRaises(ValueError):
+            trainer = CenterLossTrainer(decoder=decoder,
+                                        classifier=classifier,
+                                        feature_dim=64,
+                                        accelerator='cpu',
+                                        num_classes=2,
+                                        metrics=['unexpected'])
+            trainer.fit(train_loader, val_loader)
+            trainer.test(test_loader)
+
+        trainer = CenterLossTrainer(
+            decoder=decoder,
+            classifier=classifier,
+            feature_dim=64,
+            devices=1,
+            accelerator='gpu',
+            num_classes=2,
+            metrics=['accuracy', 'recall', 'precision', 'f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/trainers/test_imbalance_trainer.py b/test/trainers/test_imbalance_trainer.py
new file mode 100644
index 0000000..2332571
--- /dev/null
+++ b/test/trainers/test_imbalance_trainer.py
@@ -0,0 +1,297 @@
+import random
+import unittest
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset
+from torch.utils.data.dataloader import DataLoader
+from torcheeg.trainers.imbalance import LALossTrainer, LDAMLossTrainer, EQLossTrainer, FocalLossTrainer, WCELossTrainer
+
+
+class DummyDataset(Dataset):
+
+    def __init__(self, length: int = 101):
+        self.length = length
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, i) -> int:
+        return torch.randn(120), random.randint(0, 1)
+
+
+class DummyModel(nn.Module):
+
+    def __init__(self, in_channels=120, out_channels=2):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.fc = nn.Linear(in_channels, out_channels)
+
+    def forward(self, x):
+        return self.fc(x)
+
+
+class TestImbalanceTrainer(unittest.TestCase):
+
+    def test_la_loss_trainer(self):
+        train_dataset = DummyDataset()
+        val_dataset = DummyDataset()
+        test_dataset = DummyDataset()
+
+        train_loader = DataLoader(train_dataset, batch_size=1)
+        val_loader = DataLoader(val_dataset, batch_size=1)
+        test_loader = DataLoader(test_dataset, batch_size=1)
+
+        model = DummyModel()
+
+        trainer = LALossTrainer(
+            model,
+            devices=1,
+            class_frequency=train_loader,
+            accelerator='cpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+        trainer = LALossTrainer(
+            model,
+            class_frequency=[10, 20],
+            devices=1,
+            accelerator='cpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+        trainer = LALossTrainer(
+            model,
+            devices=1,
+            class_frequency=train_loader,
+            accelerator='gpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+        trainer = LALossTrainer(
+            model,
+            class_frequency=[10, 20],
+            devices=1,
+            accelerator='gpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+    def test_ldam_loss_trainer(self):
+        train_dataset = DummyDataset()
+        val_dataset = DummyDataset()
+        test_dataset = DummyDataset()
+
+        train_loader = DataLoader(train_dataset, batch_size=1)
+        val_loader = DataLoader(val_dataset, batch_size=1)
+        test_loader = DataLoader(test_dataset, batch_size=1)
+
+        model = DummyModel()
+
+        trainer = LDAMLossTrainer(
+            model,
+            devices=1,
+            class_frequency=train_loader,
+            accelerator='cpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+        trainer = LDAMLossTrainer(
+            model,
+            rule='drw',
+            class_frequency=train_loader,
+            drw_epochs=1,
+            devices=1,
+            accelerator='gpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+        trainer = LDAMLossTrainer(
+            model,
+            devices=1,
+            rule='reweight',
+            class_frequency=train_loader,
+            accelerator='gpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+        trainer = LDAMLossTrainer(
+            model,
+            class_frequency=[10, 20],
+            devices=1,
+            accelerator='gpu',
+            rule='drw',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+    def test_eq_loss_trainer(self):
+        train_dataset = DummyDataset()
+        val_dataset = DummyDataset()
+        test_dataset = DummyDataset()
+
+        train_loader = DataLoader(train_dataset, batch_size=1)
+        val_loader = DataLoader(val_dataset, batch_size=1)
+        test_loader = DataLoader(test_dataset, batch_size=1)
+
+        model = DummyModel()
+
+        trainer = EQLossTrainer(
+            model,
+            devices=1,
+            class_frequency=train_loader,
+            accelerator='cpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+        trainer = EQLossTrainer(
+            model,
+            class_frequency=[10, 20],
+            devices=1,
+            accelerator='cpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+        trainer = EQLossTrainer(
+            model,
+            devices=1,
+            class_frequency=train_loader,
+            accelerator='gpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+        trainer = EQLossTrainer(
+            model,
+            class_frequency=[10, 20],
+            devices=1,
+            accelerator='gpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+    def test_focal_loss_trainer(self):
+        train_dataset = DummyDataset()
+        val_dataset = DummyDataset()
+        test_dataset = DummyDataset()
+
+        train_loader = DataLoader(train_dataset, batch_size=1)
+        val_loader = DataLoader(val_dataset, batch_size=1)
+        test_loader = DataLoader(test_dataset, batch_size=1)
+
+        model = DummyModel()
+
+        trainer = FocalLossTrainer(
+            model,
+            devices=1,
+            class_frequency=train_loader,
+            accelerator='cpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+        trainer = FocalLossTrainer(
+            model,
+            class_frequency=[10, 20],
+            devices=1,
+            accelerator='cpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+        trainer = FocalLossTrainer(
+            model,
+            devices=1,
+            class_frequency=train_loader,
+            accelerator='gpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+        trainer = FocalLossTrainer(
+            model,
+            class_frequency=[10, 20],
+            devices=1,
+            accelerator='gpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+    def test_wce_loss_trainer(self):
+        train_dataset = DummyDataset()
+        val_dataset = DummyDataset()
+        test_dataset = DummyDataset()
+
+        train_loader = DataLoader(train_dataset, batch_size=1)
+        val_loader = DataLoader(val_dataset, batch_size=1)
+        test_loader = DataLoader(test_dataset, batch_size=1)
+
+        model = DummyModel()
+
+        trainer = WCELossTrainer(
+            model,
+            devices=1,
+            class_frequency=train_loader,
+            accelerator='cpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+        trainer = WCELossTrainer(
+            model,
+            class_frequency=[10, 20],
+            devices=1,
+            accelerator='cpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+        trainer = WCELossTrainer(
+            model,
+            devices=1,
+            class_frequency=train_loader,
+            accelerator='gpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+        trainer = WCELossTrainer(
+            model,
+            class_frequency=[10, 20],
+            devices=1,
+            accelerator='gpu',
+            num_classes=2,
+            metrics=['f1score'])
+        trainer.fit(train_loader, val_loader, max_epochs=2)
+        trainer.test(test_loader)
+
+if __name__ == '__main__':
+    unittest.main()