Thank you @ptrblck!
In fact, I am trying an attention mechanism and I want the weights that I update, according to attention, to take effect during training. Thus, from my understanding, torch.no_grad() should only be used during testing/validation. Here’s the code I have in mind (for a simple demonstration, I am replacing the attention weights here with some random values):
    model.train(); total_loss = 0
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device) 
        optimizer.zero_grad()
        attention_matrix = torch.rand(model.fc.weight.size())
        model.fc.weight = model.fc.weight * attention_matrix
        output = model(data)           
        loss = criterion(output, target)
        total_loss += loss.item()
        loss.backward()
        optimizer.step()

During testing, however, I think we need to use torch.no_grad() , as follows:
    model.eval()  
    with torch.no_grad():
        for data, target  in test_loader:
            data, target = data.to(device), target.to(device)
            model.fc.weight = model.fc.weight * attention_matrix
            output = model(data)
            loss = criterion(output, target)
            total_loss += loss.item()

I also need to do the same for the convolution layers, not yet sure how this could be done!