Browse Source

更新 'gpu_mnist_example/train_gpu.py'

liuzx-patch-1
liuzxtest02 2 years ago
parent
commit
c833765afb
1 changed files with 4 additions and 0 deletions
  1. +4
    -0
      gpu_mnist_example/train_gpu.py

+ 4
- 0
gpu_mnist_example/train_gpu.py View File

@@ -91,6 +91,10 @@ if __name__ == '__main__':
print(os.listdir(pretrain_model_path))
os.listdir(pretrain_model_path)

print("output_path:")
print(os.listdir(output_path))
os.listdir(output_path)

#log output
print('cuda is available:{}'.format(torch.cuda.is_available()))
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


Loading…
Cancel
Save