pytorch

1

2

3

4

5
from caffe2.python.normalizer_context import UseNormalizer, NormalizerContext
6
from caffe2.python.normalizer import BatchNormalizer
7
from caffe2.python.layer_test_util import LayersTestCase
8

9

10
class TestNormalizerContext(LayersTestCase):
11
    def test_normalizer_context(self):
12
        bn = BatchNormalizer(momentum=0.1)
13
        with UseNormalizer({'BATCH': bn}):
14
            normalizer = NormalizerContext.current().get_normalizer('BATCH')
15
            self.assertEqual(bn, normalizer)
16