diff --git a/mediapipe/model_maker/python/core/hyperparameters.py b/mediapipe/model_maker/python/core/hyperparameters.py
index 22471655..92e1856c 100644
--- a/mediapipe/model_maker/python/core/hyperparameters.py
+++ b/mediapipe/model_maker/python/core/hyperparameters.py
@@ -15,7 +15,7 @@
 
 import dataclasses
 import tempfile
-from typing import Optional
+from typing import Mapping, Optional
 
 import tensorflow as tf
 
@@ -36,6 +36,8 @@ class BaseHParams:
     steps_per_epoch: An optional integer indicate the number of training steps
       per epoch. If not set, the training pipeline calculates the default steps
       per epoch as the training dataset size divided by batch size.
+    class_weights: An optional mapping of indices to weights for weighting the
+      loss function during training.
     shuffle: True if the dataset is shuffled before training.
     export_dir: The location of the model checkpoint files.
     distribution_strategy: A string specifying which Distribution Strategy to
@@ -57,6 +59,7 @@ class BaseHParams:
   batch_size: int
   epochs: int
   steps_per_epoch: Optional[int] = None
+  class_weights: Optional[Mapping[int, float]] = None
 
   # Dataset-related parameters
   shuffle: bool = False
diff --git a/mediapipe/model_maker/python/core/tasks/classifier.py b/mediapipe/model_maker/python/core/tasks/classifier.py
index a042c0ec..d504defb 100644
--- a/mediapipe/model_maker/python/core/tasks/classifier.py
+++ b/mediapipe/model_maker/python/core/tasks/classifier.py
@@ -110,7 +110,9 @@ class Classifier(custom_model.CustomModel):
         # dataset is exhausted even if there are epochs remaining.
         steps_per_epoch=None,
         validation_data=validation_dataset,
-        callbacks=self._callbacks)
+        callbacks=self._callbacks,
+        class_weight=self._hparams.class_weights,
+    )
 
   def evaluate(self, data: dataset.Dataset, batch_size: int = 32) -> Any:
     """Evaluates the classifier with the provided evaluation dataset.