Based on the code it seems you are currently using 1 conv, 2 linear, 1 pool, 2 relu modules as well as your custom criterion, so I would start with a single conv layer and make sure the outputs as well as the gradients match. Once this is done, add layers one by one and make sure the test is still working until you are able to verify the entire block.