Code change:
Hi, I need to pass input through one nn.Module, then argmax the output, and then pass that output to a second nn.Module. I want to backprop through the argmax back to the weights of the first module. That's impossible right now, though, it seems. When I loop through the gradient of the weights of my first module (with the following code): firstmodel.zerograd() secondmodel.zerograd() loss.backward() # retainvariables=True for param in secondmodel.parameters(): print( 'data: ', ) print( 'grad data: ', ) for param in firstmodel.parameters(): print( 'data: ', ) print( 'grad data: ', ) I get the following error, stating that those gradients can no longer be viewed/exist: -main()
argmax is not differentiable (and even if you try to think about it as a function differentiable almost everywhere, its derivative is always 0, so you won't get anything meaningful).

Related discussion:
text: The transposed convolution operator multiplies each input value element-wise by a learnable kernel, and sums over the outputs from all input feature planes. <del> **This module can be seen as the exact reverse of Conv3d**. It is sometimes (but incorrectly) refered to as a deconvolutional operation. </del> <ins> This module can be seen as the gradient of Conv3d with respect to its input. It is also known as a fractionally-strided convolution or a deconvolution (although it is not an actual deconvolution operation). </ins> | :attr:`stride` controls the stride for the cross-correlation. | If :attr:`padding` is non-zero, then the input is implicitly zero-padded on both sides <del> for :attr:`padding` number of points </del> <ins> for :attr:`padding` number of points. </ins> | If :attr:`output_padding` is non-zero, then the output is implicitly zero-padded on one side <del> for :attr:`output_padding` number of points | :attr:`groups` controls the connections between inputs and outputs. </del> <ins> for :attr:`output_padding` number of points. | :attr:`dilation` controls the spacing between the kernel points; also known as the à trous algorithm. It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does. | :attr:`groups` controls the connections between inputs and outputs. `in_channels` and `out_channels` must both be divisible by `groups`. </ins> |       At groups=1, all inputs are convolved to all outputs. |       At groups=2, the operation becomes equivalent to having two conv layers side by side, each seeing half the input channels, and producing half the output channels, and both subsequently concatenated. <ins> At groups=`in_channels`, each input channel is convolved with its own set of filters (of size `out_channels // in_channels`). </ins> The parameters :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, :attr:`output_padding` can either be: <del> - a single ``int`` -- in which case the same value is used for the height and width dimension </del> <ins> - a single ``int`` -- in which case the same value is used for the depth, height and width dimensions </ins> - a ``tuple`` of three ints -- in which case, the first `int` is used for the depth dimension, the second `int` for the width dimension and the third `int` for the width dimension