pytorch/caffe2/python/normalizer_test.py
Aaron Gokaslan 748bac8757 [BE]: Apply pyupgrade yield from and unit test alias upgrades (#94309)
Applies some more harmless pyupgrades. This one gets rid of deprecated aliases in unit_tests and more upgrades yield for loops into yield from generators which are more performance and propagates more information / exceptions from original generator. This is the modern recommended way of forwarding generators.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/94309
Approved by: https://github.com/albanD
2023-02-07 20:08:58 +00:00

16 lines
486 B
Python

from caffe2.python.normalizer_context import UseNormalizer, NormalizerContext
from caffe2.python.normalizer import BatchNormalizer
from caffe2.python.layer_test_util import LayersTestCase
class TestNormalizerContext(LayersTestCase):
def test_normalizer_context(self):
bn = BatchNormalizer(momentum=0.1)
with UseNormalizer({'BATCH': bn}):
normalizer = NormalizerContext.current().get_normalizer('BATCH')
self.assertEqual(bn, normalizer)