not sure about backward pass, but forward pass is under control. let's train a whole net like this and in the meantime see if backward pass can be improved / is correct