I am trying to update a particular mask based on the values shown in each sample. So far, I can update the mask based on the entire batch, but it only updates after all the samples have passed through.
Is there some way I can get in there and call an event or make something happen after each sample in a batch as opposed to the whole batch?
Example pseudo code.
Let’s assume x is a tensor of size [32, 1000]
def forward(self, x):
    x = F.relu(self.fc1(x))
    x = x * mask
    x = F.sigmoid(x)
    mask = mask.update(x)
    return x

This code will successfully update, and the next batch around It will apply the updated mask, but I’m trying to get it to update after every sample, not after the entire batch.