error when using rmse loss function with tabnet regressor

I want to use RMSE as a loss function for regression problems but when I try to do it, it gives me an error. below is the code

def loss_fn(x,y):
    criterion = nn.MSELoss()
    return torch.sqrt(criterion(x, y))

clf1 = TabNetRegressor(optimizer_fn=torch.optim.Adam,
                       cat_idxs=[0,1,2,3,4,5,6,7,8,9], cat_dims=[2,2,2,4,4,4,8,8,7,15],lambda_sparse=0,
                       optimizer_params=dict(lr=2e-2),
                       scheduler_params={"step_size":10, "gamma":0.9},
                       scheduler_fn=torch.optim.lr_scheduler.StepLR, mask_type='entmax')

# fit the model 
clf1.fit(
    x_train,y_train,
    eval_set=[(x_train, y_train), (x_val, y_val)],loss_fn=loss_fn,
    eval_name=['train', 'valid'],
    eval_metric=['mse','rmse'],
    max_epochs=50 , patience=100,
    batch_size=100000, virtual_batch_size=50000,
    
)```

below is the error 



![Screenshot from 2021-08-27 16-11-28](https://user-images.githubusercontent.com/81755307/131115420-e6748758-3e70-414b-be3b-493d10da527c.png)

[Optimox]

well I guess it’s totally ok to switch with loss = loss - something but I’d be curious to understand what’s the problem behind that.

[YashVardhan-AI]

i did some testing and its the problem with torch.sqrt other techniques work fine