1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
|
# Owner(s): ["oncall: mobile"]
import torch
from torch.testing._internal.common_utils import TestCase, run_tests
class TestSetDefaultMobileCPUAllocator(TestCase):
def test_no_exception(self):
torch._C._set_default_mobile_cpu_allocator()
torch._C._unset_default_mobile_cpu_allocator()
def test_exception(self):
with self.assertRaises(Exception):
torch._C._unset_default_mobile_cpu_allocator()
with self.assertRaises(Exception):
torch._C._set_default_mobile_cpu_allocator()
torch._C._set_default_mobile_cpu_allocator()
# Must reset to good state
# For next test.
torch._C._unset_default_mobile_cpu_allocator()
with self.assertRaises(Exception):
torch._C._set_default_mobile_cpu_allocator()
torch._C._unset_default_mobile_cpu_allocator()
torch._C._unset_default_mobile_cpu_allocator()
if __name__ == '__main__':
run_tests()
|
