python-pytorch-rocm/trunk/test.py

   1 #!/usr/bin/env python
   2
   3 import torch
   4
   5 if torch.cuda.is_available() and torch.version.hip:
   6     d = torch.device('cuda')
   7     a = torch.rand(1, 2).to(d)
   8     print(a + 0)
   9     print("TESTS PASSED!")
  10 else:
  11     print("ROCm not supported with PyTorch installation")