mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
b004307252
Test Plan: manual inspection & sandcastle Reviewed By: zertosh Differential Revision: D30279364 fbshipit-source-id: c1ed77dfe43a3bde358f92737cd5535ae5d13c9a
30 lines
946 B
Python
30 lines
946 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()
|