How to properly prepare a validation set for the SAITS model

[TDL77]

I have a dataset. containing many Nan. size X_3d = (12352, 1080, 96) , I also made a validation dataset from it, trying to remove all the Nan and leave it. but it still falls into error when I use the validation sample.
X_intact, X_val, missing_mask, indicating_mask = mcar(X_val, 0.2) # hold out 20% observed values as ground truth
X_val = masked_fill(X_val, 1 - missing_mask, np.nan)
, it has the size (4399, 1080, 96)

dataset = {"X" : X_3d}
val_data = {'X': X_val}
#train the model. Here I use the whole dataset as the training set, because ground truth is not visible to the model.
saits.fit(train_set= dataset , val_set=val_data)

---

KeyError Traceback (most recent call last)
File c:\Users\admin\mambaforge\envs\pypots\lib\site-packages\pypots\imputation\base.py:294, in BaseNNImputer._train_model(self, training_loader, val_loader)
290 imputation_collector = imputation_collector.cpu().detach().numpy()
292 mean_val_loss = cal_mae(
293 imputation_collector,
--> 294 val_loader.dataset.data["X_intact"],
295 val_loader.dataset.data["indicating_mask"],
296 # the above val_loader.dataset.data is a dict containing the validation dataset
297 )
299 # save validating loss logs into the tensorboard file for every epoch if in need

KeyError: 'X_intact'

During handling of the above exception, another exception occurred:

RuntimeError Traceback (most recent call last)
Cell In[93], line 24
22 val_data = {'X': X_val}
23 # train the model. Here I use the whole dataset as the training set, because ground truth is not visible to the model.
---> 24 saits.fit(train_set= dataset , val_set=val_data)
25 imputation = saits.impute(dataset) # impute the originally-missing values and artificially-missing values

File c:\Users\admin\mambaforge\envs\pypots\lib\site-packages\pypots\imputation\saits\model.py:420, in SAITS.fit(self, train_set, val_set, file_type)
412 val_loader = DataLoader(
413 val_set,
414 batch_size=self.batch_size,
415 shuffle=False,
416 num_workers=self.num_workers,
417 )
419 # Step 2: train the model and freeze it
--> 420 self._train_model(training_loader, val_loader)
421 self.model.load_state_dict(self.best_model_dict)
422 self.model.eval() # set the model as eval status to freeze it.

File c:\Users\admin\mambaforge\envs\pypots\lib\site-packages\pypots\imputation\base.py:342, in BaseNNImputer._train_model(self, training_loader, val_loader)
340 logger.error(f"Exception: {e}")
341 if self.best_model_dict is None:
--> 342 raise RuntimeError(
343 "Training got interrupted. Model was not trained. Please investigate the error printed above."
344 )
345 else:
346 RuntimeWarning(
347 "Training got interrupted. Please investigate the error printed above.\n"
348 "Model got trained and will load the best checkpoint so far for testing.\n"
349 "If you don't want it, please try fit() again."
350 )

RuntimeError: Training got interrupted. Model was not trained. Please investigate the error printed above.