@ptrblck, Thankyou for your response.
So the mask_compute function takes the output of net1 (size = Nx1) and creates a mask of size (Nx1) where the elements of the mask are 1 if the value of the corresponding element of the output of net1 is above a threshold and 0 if it’s below a threshold.
I am not sure how the forward and backward pass can be implemented for this?