mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
a4a6d056e6
Summary: Fixes https://github.com/pytorch/pytorch/issues/66232 This should be the last immediate task. I anticipate test ownership will change overtime but this is the last big thing to close it out Pull Request resolved: https://github.com/pytorch/pytorch/pull/67859 Reviewed By: soulitzer Differential Revision: D32210534 Pulled By: janeyx99 fbshipit-source-id: 7fd835d87d9d35d49ec49de1fcfa29b085133e99
30 lines
976 B
Python
30 lines
976 B
Python
# 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()
|