Due to certain data I want my network to only learn on losses that are lower than 1. For that I use the following code:
def train(img, label):
   x = self.net(img)
    loss_val = self.loss(x, label)
    self.optimizer.zero_grad()
        
    if loss_val < 1:
        loss_val.backward()
        self.optimizer.step()

Is this correct? Because I have the feeling that the gradients are not reseting so as soon as a loss_val is under 1, all updates on the gradients are done, even the ones on which the osses were higher than 1.
How can I otherwise reset all the gradients from the optimiser?