So, I corrected the loss function. The error was in the script. Earlier, I was enabling the gradients inside each training batch as follows:
for epoch in range(params.train.num_epochs):
    for train_batch, train_data in enumerate(trainloader):
        model.train()
        optimizer.zero_grad()
        inputs = torch.autograd.Variable(train_data[0], requires_grad=False).to(device)
        labels = torch.autograd.Variable(train_data[1], requires_grad=False).to(device)
        with torch.set_grad_enabled(True):
            outputs = model(inputs)
            loss = torch.sqrt(loss_criterion(outputs,labels))
            loss.backward()
            optimizer.step()

So, this time, I enabled them outside the loop of training batches:
for epoch in range(params.train.num_epochs):
    torch.set_grad_enabled(True)
    for train_batch, train_data in enumerate(trainloader):
        model.train()
        optimizer.zero_grad()
        inputs = torch.autograd.Variable(train_data[0], requires_grad=False).to(device)
        labels = torch.autograd.Variable(train_data[1], requires_grad=False).to(device)
        outputs = model(inputs)
        loss = torch.sqrt(loss_criterion(outputs,labels))
        loss.backward()
        optimizer.step()

And it worked!