class MyLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.mse = nn.MSELoss()

    def forward(self, yhat, y):
        loss = 2*self.mse(yhat[:, :2], y[:, :2]) + 5*self.mse(yhat[:, 2:5], y[:, 2:5]) + self.mse(yhat[:, 5:], y[:, 5:])
        return loss