1 # Copyright 2016 The TensorFlow Authors. All Rights Reserved. 2 # 3 # Licensed under the Apache License, Version 2.0 (the "License"); 4 # you may not use this file except in compliance with the License. 5 # You may obtain a copy of the License at 6 # 7 # http://www.apache.org/licenses/LICENSE-2.0 8 # 9 # Unless required by applicable law or agreed to in writing, software 10 # distributed under the License is distributed on an "AS IS" BASIS, 11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 # See the License for the specific language governing permissions and 13 # limitations under the License. 14 # ============================================================================== 15 """Tests for the SWIG-wrapped device lib.""" 16 17 from __future__ import absolute_import 18 from __future__ import division 19 from __future__ import print_function 20 21 from tensorflow.python.client import device_lib 22 from tensorflow.python.framework import test_util 23 from tensorflow.python.platform import googletest 24 from tensorflow.python.platform import test 25 26 27 class DeviceLibTest(test_util.TensorFlowTestCase): 28 29 def testListLocalDevices(self): 30 devices = device_lib.list_local_devices() 31 self.assertGreater(len(devices), 0) 32 self.assertEqual(devices[0].device_type, "CPU") 33 34 # GPU test 35 if test.is_gpu_available(): 36 self.assertGreater(len(devices), 1) 37 self.assertTrue("GPU" in [d.device_type for d in devices] or 38 "SYCL" in [d.device_type for d in devices]) 39 40 41 if __name__ == "__main__": 42 googletest.main() 43
