pytorch

1
# Owner(s): ["oncall: mobile"]
2

3
import torch
4
from torch.testing._internal.common_utils import TestCase, run_tests
5

6
class TestSetDefaultMobileCPUAllocator(TestCase):
7
    def test_no_exception(self):
8
        torch._C._set_default_mobile_cpu_allocator()
9
        torch._C._unset_default_mobile_cpu_allocator()
10

11
    def test_exception(self):
12
        with self.assertRaises(Exception):
13
            torch._C._unset_default_mobile_cpu_allocator()
14

15
        with self.assertRaises(Exception):
16
            torch._C._set_default_mobile_cpu_allocator()
17
            torch._C._set_default_mobile_cpu_allocator()
18

19
        # Must reset to good state
20
        # For next test.
21
        torch._C._unset_default_mobile_cpu_allocator()
22

23
        with self.assertRaises(Exception):
24
            torch._C._set_default_mobile_cpu_allocator()
25
            torch._C._unset_default_mobile_cpu_allocator()
26
            torch._C._unset_default_mobile_cpu_allocator()
27

28
if __name__ == '__main__':
29
    run_tests()
30