--- a
+++ b/tests/test_sampler.py
@@ -0,0 +1,78 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import pytest
+import torch
+
+from mmseg.core import OHEMPixelSampler
+from mmseg.models.decode_heads import FCNHead
+
+
+def _context_for_ohem():
+    return FCNHead(in_channels=32, channels=16, num_classes=19)
+
+
+def _context_for_ohem_multiple_loss():
+    return FCNHead(
+        in_channels=32,
+        channels=16,
+        num_classes=19,
+        loss_decode=[
+            dict(type='CrossEntropyLoss', loss_name='loss_1'),
+            dict(type='CrossEntropyLoss', loss_name='loss_2')
+        ])
+
+
+def test_ohem_sampler():
+
+    with pytest.raises(AssertionError):
+        # seg_logit and seg_label must be of the same size
+        sampler = OHEMPixelSampler(context=_context_for_ohem())
+        seg_logit = torch.randn(1, 19, 45, 45)
+        seg_label = torch.randint(0, 19, size=(1, 1, 89, 89))
+        sampler.sample(seg_logit, seg_label)
+
+    # test with thresh
+    sampler = OHEMPixelSampler(
+        context=_context_for_ohem(), thresh=0.7, min_kept=200)
+    seg_logit = torch.randn(1, 19, 45, 45)
+    seg_label = torch.randint(0, 19, size=(1, 1, 45, 45))
+    seg_weight = sampler.sample(seg_logit, seg_label)
+    assert seg_weight.shape[0] == seg_logit.shape[0]
+    assert seg_weight.shape[1:] == seg_logit.shape[2:]
+    assert seg_weight.sum() > 200
+
+    # test w.o thresh
+    sampler = OHEMPixelSampler(context=_context_for_ohem(), min_kept=200)
+    seg_logit = torch.randn(1, 19, 45, 45)
+    seg_label = torch.randint(0, 19, size=(1, 1, 45, 45))
+    seg_weight = sampler.sample(seg_logit, seg_label)
+    assert seg_weight.shape[0] == seg_logit.shape[0]
+    assert seg_weight.shape[1:] == seg_logit.shape[2:]
+    assert seg_weight.sum() == 200
+
+    # test multiple losses case
+    with pytest.raises(AssertionError):
+        # seg_logit and seg_label must be of the same size
+        sampler = OHEMPixelSampler(context=_context_for_ohem_multiple_loss())
+        seg_logit = torch.randn(1, 19, 45, 45)
+        seg_label = torch.randint(0, 19, size=(1, 1, 89, 89))
+        sampler.sample(seg_logit, seg_label)
+
+    # test with thresh in multiple losses case
+    sampler = OHEMPixelSampler(
+        context=_context_for_ohem_multiple_loss(), thresh=0.7, min_kept=200)
+    seg_logit = torch.randn(1, 19, 45, 45)
+    seg_label = torch.randint(0, 19, size=(1, 1, 45, 45))
+    seg_weight = sampler.sample(seg_logit, seg_label)
+    assert seg_weight.shape[0] == seg_logit.shape[0]
+    assert seg_weight.shape[1:] == seg_logit.shape[2:]
+    assert seg_weight.sum() > 200
+
+    # test w.o thresh in multiple losses case
+    sampler = OHEMPixelSampler(
+        context=_context_for_ohem_multiple_loss(), min_kept=200)
+    seg_logit = torch.randn(1, 19, 45, 45)
+    seg_label = torch.randint(0, 19, size=(1, 1, 45, 45))
+    seg_weight = sampler.sample(seg_logit, seg_label)
+    assert seg_weight.shape[0] == seg_logit.shape[0]
+    assert seg_weight.shape[1:] == seg_logit.shape[2:]
+    assert seg_weight.sum() == 200