Thank you so much for this. It answered my first question so well. Let me consolidate my code and ask the same question better.
My next and last question is how I can undo loss.backward() for a given batch. I wish to compute the autograd backprop gradients. But, I do not wish to apply those gradients to model computation graph. How could I do that?
The final objective is to get two sets of gradients for each batch (regular and noisy one) and apply the gradient set that produces lower loss values. How can I achieve that?
The cleaner code (thanks to @albanD ) is as follows:
def train_epoch(eta, model, train_loader, criterion):
    model.train()
    running_loss = 0.0; predictions = []; ground_truth = []; loss_den = 1
    
    start_time = time.time()
    optimiser = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
    
    for batch_idx, (data, target) in enumerate(train_loader):
        
        data = data.to(device); target = target.to(device)
    
        #previous model
        outputs = model(data.float())        
        loss = criterion(outputs, target)
        loss.backward()
        optimiser.step()
        
        #convGrad is the set of old gradients
        oldGrad = list()
        for param in model.parameters():
            convGrad = param.grad
            oldGrad.append(convGrad)
            noisyGrad = eta * np.abs(convGrad.detach().cpu().numpy())
            newGrad = convGrad + torch.from_numpy(np.random.uniform(-noisyGrad, noisyGrad))
            param.grad = nn.Parameter(torch.from_numpy(newGrad.detach().numpy()).float())
        
        #The new loss value for the new gradients is computed
        outputs = model(data.float())
        loss_new = criterion(outputs, target)
        loss_den += 1

        if loss_new.item() > loss.item():
            for paramIdx, param in enumerate(model.parameters()):
                param.grad = oldGrad[paramIdx]

            running_loss += loss.item()
        else:
            running_loss += loss_new.item()
        
    end_time = time.time()
    running_loss /= loss_den
    print('Training Loss: ', running_loss, 'Time: ',end_time - start_time, 's')
    
    return running_loss, model