flora:

First of all, I have a question regarding the execution speed. MyConv2d is much slower than conv2d (because of the double for loop I guess). Is there a way to speed it up?


You could try to remove the loops and unfold the data, which could use more memory but might also be faster. Alternatively, you could also write a custom C++/CUDA extension, which could also yield a speedup.



 flora:

Secondly, I have an issue at the very first iteration when I train my network on gpu. Indeed, once the input got through my first custom layer, I get back Nan values in the output. Do you have any idea why this happens? Is there something wrong with my implementation of MyConv2d?


You could add debug print statements and check which part of your custom layer returns the NaN values to narrow it down further.



 flora:

Last, I recently have a weird error that came out of the blue when I train my network:


If you are using an older PyTorch version, please update to the latest stable, since indexing errors should raise RuntimeErrors, not fail with illegal memory accesses.