[PYTHON] Added option to show PTX source code in Python

2020-11-07 02:55:48 -05:00
parent cf80ccc798
commit 8f3ee53f24
10 changed files with 68 additions and 33 deletions
--- a/python/examples/tutorials/mat_mul.py
+++ b/python/examples/tutorials/mat_mul.py
@@ -121,8 +121,8 @@ dot = _dot.apply
 torch.manual_seed(0)

 M, N, K = 2048, 2048, 2048
-a = torch.rand((M, K)).cuda()
-b = torch.rand((K, N)).cuda()
+a = torch.rand((M, K)).cuda().half()
+b = torch.rand((K, N)).cuda().half()

 #a[:] = 1
 #b[:] = 1