OK, I ended up changing the loss function as follows, How ever is it correct? or did I mess something up here again?
def loss_function(output_e, outputs, imgs, device):
    criterion = nn.MSELoss()
    assert outputs.shape == imgs.shape ,f'outputs.shape : {outputs.shape} != imgs.shape : {imgs.shape}'
    
    loss1 = criterion(outputs, imgs)
    output_e.backward(torch.ones(output_e.size()).to(device), retain_graph=True)    
    loss2 = torch.mean(pow(imgs.grad,2))
    imgs.grad.data.zero_()
    loss = loss1 + loss2 
    return loss 

the notable changes are


instead ofimgs.grad.requires_grad=True, I used imgs.requires_grad=True as it doesn’t make any sense to me, to make grads require grads! while imgs doesn’t require it in first place and it will always result in error (NoneType doesnt have require_grads!)


in order for imgs to have gradients, you need to remember:
First imgs is a non-leaf node. autograd wont store grads for non-leaf nodes. in order to make them have gradients, you should use imgs.retain_grad().  retain_grad() must be called before doing forward().
Second.requires_grad is not retroactive, which means it must be set prior to running forward()


So the training loop should look like this ultimately :
for e in range(epochs):
    for i, (imgs, labels) in enumerate(dataloader_train):
        imgs = imgs.to(device)
        labels = labels.to(device)

        imgs.retain_grad()
        imgs.requires_grad_(True)
        
        outputs_e, outputs = model(imgs)
        loss = loss_function(outputs_e, outputs, imgs, lam,device)

        imgs.requires_grad_(False)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print(f'epoch/epochs: {e}/{epochs} loss: {loss.item():.4f}')