I wrote a custom CUDA module, following the provided example: https://github.com/pytorch/extension-ffi 11
So far, my modules worked flawlessly. Note that unlike the example, I am using custom CUDA kernels. However, I just wanted to start training my model on multiple GPUs and run into ‘illegal memory access’ exceptions in the backward call. Since the custom CUDA module is working perfectly fine on a single GPU and since the exact same CUDA code is working perfectly fine with Torch on multiple GPUs, I suspect that something in my PyTorch related Python code is wrong:
class CustomFunction(torch.autograd.Function):
	def __init__(self):
		super(CustomFunction, self).__init__()
	# end

	def forward(self, input1, input2, input3):
		self.save_for_backward(input1, input2, input3)

		assert(input1.is_contiguous() == True)
		assert(input2.is_contiguous() == True)
		assert(input3.is_contiguous() == True)

		output = input1.new(...).zero_()

		# call C portion which call the CUDA portion whcih calls the forward kernel

		return output
	# end

	def backward(self, gradOutput):
		input1, input2, input3 = self.saved_tensors

		assert(gradOutput.is_contiguous() == True)

		gradInput1 = input1.new(...).zero_()
		gradInput2 = input1.new(...).zero_()
		gradInput3 = input1.new(...).zero_()

		# call C portion which call the CUDA portion whcih calls the backward kernel

		return gradInput1, gradInput2, gradInput3
	# end
# end

When the call to the CUDA kernel in the backward function is removed, the ‘illegal memory access’ error does not appear but the model does of course not train correctly due to the missing gradients. I thus suspect some issue with the memory allocation and have already tried using gradOutput.new instead of input1.new in the backward function. I likewise used .get_device() to make sure that all the tensors in a call of the backward function are on the same GPU and can confirm that they are.
I normally am very reluctant to ask for help but am afraid that I might be missing something fundamental here. Is anyone able to provide any insight here? Thank you very much for making PyTorch happen by the way, it so far has been a joy to work with!