Add architecture debugging

This commit is contained in:
Brandon Thomas
2022-06-10 17:09:33 -04:00
parent 44f2a1082b
commit 711d5adb73
+2
View File
@@ -35,6 +35,8 @@ def print_gpu_info():
print('PyTorch version', torch.__version__)
print('CUDA Available?', torch.cuda.is_available())
print('CUDA Device count', torch.cuda.device_count())
print('CUDA architectures library was compiled for', torch.cuda.get_arch_list())
#try:
# from tensorflow.python.client import device_lib
# print('local devices', str(device_lib.list_local_devices()).replace("\n", "\n "))