import torch
import torch.nn as nn


class RMSELoss(nn.Module):
  def __init__(self, eps=1e-7):
    super().__init__()
    self.mse = nn.MSELoss()
    self.eps = eps
  
  def forward(self, yhat, y):
    loss = torch.sqrt(self.mse(yhat, y) + self.eps)
    return loss
