Hello all! In my current project, I have pass in 2D array of values (like [4096, 3] or [32768, 3]) into a neural network which processes these values and returns just [4096] or [32768]. The problem with this, as one might imagine, is that you run out of memory super fast at higher resolutions of input (like [64x64x64, 3]).

Luckily, only part of the values calculated from the neural network are necessary for calculations. Most of the values returned will be negative, some will be positive, and I only need to perform calculations on the positive values. Therefore, all those negative values just end up being unnecessary calculations for backprop and take up memory. Therefore, I was wondering if it’s possible to remove the extra calculations from the computational graph so that I can still work with higher resolutions and not have to backprop unneeded data.