When we use nn.ReLU(inplace=True), it should not improve the GPU memory for back-propagation. This is true for convolutional neural networks.
However, the inplace mode doesn’t work for custom functions.
Specifically, I build a custom autograd function and a custom module based on it.
When I combine this module and inplace relu, it seems that nn.ReLU(inplace=True) doesn’t save memory for BP. Here is an example:
net_1 = nn.Sequential()
net_2 = nn.Sequential()
for i in range(50):
    net_1.add_module("conv_{}".format(i), MyConv2d(32, 32, 3))
    net_2.add_module("conv_{}".format(i), MyConv2d(32, 32, 3))
    net_1.add_module("relu_{}".format(i), nn.ReLU(True))
    net_2.add_module("relu_{}".format(i), nn.ReLU(False))

Then I get the max_memory_allocated for net_1 and net_2 given the same input.
And I find the used memories are the same.
net_1: 470MB
net_2: 470MB
How can I modify my custom function to make nn.ReLU(inplace=True) works?