mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
1022443168
Test Plan: revert-hammer
Differential Revision:
D30279364 (b004307252)
Original commit changeset: c1ed77dfe43a
fbshipit-source-id: eab50857675c51e0088391af06ec0ecb14e2347e
28 lines
944 B
Python
28 lines
944 B
Python
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()
|