RMSE loss for pytorch (from pytorch discussion)

RMSE_loss_pytorch.py

class RMSELoss(nn.Module):
    
    def __init__(self):
        super().__init__()
        self.mse = nn.MSELoss()
        
    def forward(self, yhat, y):
        return torch.sqrt(self.mse(yhat, y))