When training, if loss is very small, I want loss not backward to save training time by this code:
total_loss = criterion(anchor_f, pos_f, neg_f)        
if total_loss.item() > 1e-6:
    optimizer.zero_grad()
    total_loss.backward()
    optimizer.step()
else:
    pass 

this will raise cuda memory error:
RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 10.76 GiB total capacity; 9.93 GiB already allocated; 22.19 MiB free; 39.84 MiB cached)

Is my code not correct?