Stephen Merity Smerity

Smerity / cupy-pytorch-ptx.py

Created May 21, 2017 23:21 — forked from szagoruyko/cupy-pytorch-ptx.py

CuPy example for PyTorch updated to support Python 3

	import torch
	from cupy.cuda import function
	from pynvrtc.compiler import Program
	from collections import namedtuple

	a = torch.randn(1,4,4).cuda()
	b = torch.zeros(a.size()).cuda()

	kernel = '''
	extern "C"