Solved by MariosOreo in post #2 
              
              
                Hello, 
You could just call .backward() without optimizer.step() if the step is unexpected condition. 
optimizer.zero_grad()
if condition:
    loss.backward()
else:
    loss.backward()
    optimizer.step()