ValueError during training of TemporalFusionTransfomer

I'm encountering a ValueError on some datasets when training the model. The error comes from the sklearn code that checks the lengths of the predicted and actual y arrays before calculating metrics. I've checked the data itself, and there are no missing values or gaps in any of the time series. The length of my test data is 7252 (consistent with what the error is printing). The time-series length is also uniform, as are the date ranges. I'm including the stacktrace below.

`---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
in ()
28
29 # Train the model
---> 30 automl.fit(
31 X_train=X_train,
32 y_train=y_train,

/local_disk0/.ephemeral_nfs/envs/pythonEnv-d550cdcd-9d6a-4e78-ab48-3cbb09cf108a/lib/python3.9/site-packages/flaml/automl.py in fit(self, X_train, y_train, dataframe, label, metric, task, n_jobs, log_file_name, estimator_list, time_budget, max_iter, sample, ensemble, eval_method, log_type, model_history, split_ratio, n_splits, log_training_metric, mem_thres, pred_time_limit, train_time_limit, X_val, y_val, sample_weight_val, groups_val, groups, verbose, retrain_full, split_type, learner_selector, hpo_method, starting_points, seed, n_concurrent_trials, keep_search_state, preserve_checkpoint, early_stop, append_log, auto_augment, min_sample_size, use_ray, metric_constraints, custom_hp, cv_score_agg_func, skip_transform, fit_kwargs_by_estimator, **fit_kwargs)
2893 with training_log_writer(log_file_name, append_log) as save_helper:
2894 self._training_log = save_helper
-> 2895 self._search()
2896 else:
2897 self._training_log = None

/local_disk0/.ephemeral_nfs/envs/pythonEnv-d550cdcd-9d6a-4e78-ab48-3cbb09cf108a/lib/python3.9/site-packages/flaml/automl.py in _search(self)
3450 )
3451 elif self._use_ray is False:
-> 3452 self._search_sequential()
3453 else:
3454 self._search_parallel()

/local_disk0/.ephemeral_nfs/envs/pythonEnv-d550cdcd-9d6a-4e78-ab48-3cbb09cf108a/lib/python3.9/site-packages/flaml/automl.py in _search_sequential(self)
3268 )
3269 start_run_time = time.time()
-> 3270 analysis = tune.run(
3271 search_state.training_function,
3272 search_alg=search_state.search_alg,

/local_disk0/.ephemeral_nfs/envs/pythonEnv-d550cdcd-9d6a-4e78-ab48-3cbb09cf108a/lib/python3.9/site-packages/flaml/tune/tune.py in run(evaluation_function, config, low_cost_partial_config, cat_hp_cost, metric, mode, time_budget_s, points_to_evaluate, evaluated_rewards, resource_attr, min_resource, max_resource, reduction_factor, scheduler, search_alg, verbose, local_dir, num_samples, resources_per_trial, config_constraints, metric_constraints, max_failure, use_ray, use_incumbent_result_in_evaluation, log_file_name, **ray_args)
515 if verbose:
516 logger.info(f"trial {num_trials} config: {trial_to_run.config}")
--> 517 result = evaluation_function(trial_to_run.config)
518 if result is not None:
519 if isinstance(result, dict):

/local_disk0/.ephemeral_nfs/envs/pythonEnv-d550cdcd-9d6a-4e78-ab48-3cbb09cf108a/lib/python3.9/site-packages/flaml/automl.py in _compute_with_config_base(config_w_resource, state, estimator)
348 _,
349 pred_time,
--> 350 ) = compute_estimator(
351 sampled_X_train,
352 sampled_y_train,

/local_disk0/.ephemeral_nfs/envs/pythonEnv-d550cdcd-9d6a-4e78-ab48-3cbb09cf108a/lib/python3.9/site-packages/flaml/ml.py in compute_estimator(X_train, y_train, X_val, y_val, weight_val, groups_val, budget, kf, config_dic, task, estimator_name, eval_method, eval_metric, best_val_loss, n_jobs, estimator_class, cv_score_agg_func, log_training_metric, fit_kwargs)
589
590 if "holdout" == eval_method:
--> 591 val_loss, metric_for_logging, train_time, pred_time = get_val_loss(
592 config_dic,
593 estimator,

/local_disk0/.ephemeral_nfs/envs/pythonEnv-d550cdcd-9d6a-4e78-ab48-3cbb09cf108a/lib/python3.9/site-packages/flaml/ml.py in get_val_loss(config, estimator, X_train, y_train, X_val, y_val, weight_val, groups_val, eval_metric, obj, labels, budget, log_training_metric, fit_kwargs)
411 # fit_kwargs['y_val'] = y_val
412 estimator.fit(X_train, y_train, budget, **fit_kwargs)
--> 413 val_loss, metric_for_logging, pred_time, _ = _eval_estimator(
414 config,
415 estimator,

/local_disk0/.ephemeral_nfs/envs/pythonEnv-d550cdcd-9d6a-4e78-ab48-3cbb09cf108a/lib/python3.9/site-packages/flaml/ml.py in _eval_estimator(config, estimator, X_train, y_train, X_val, y_val, weight_val, groups_val, eval_metric, obj, labels, log_training_metric, fit_kwargs)
349 pred_time = (time.time() - pred_start) / X_val.shape[0]
350
--> 351 val_loss = metric_loss_score(
352 eval_metric,
353 y_processed_predict=val_pred_y,

/local_disk0/.ephemeral_nfs/envs/pythonEnv-d550cdcd-9d6a-4e78-ab48-3cbb09cf108a/lib/python3.9/site-packages/flaml/ml.py in metric_loss_score(metric_name, y_processed_predict, y_processed_true, labels, sample_weight, groups)
146 # y_processed_predict and y_processed_true are processed id labels if the original were the token labels
147 if is_in_sklearn_metric_name_set(metric_name):
--> 148 return sklearn_metric_loss_score(
149 metric_name,
150 y_processed_predict,

/local_disk0/.ephemeral_nfs/envs/pythonEnv-d550cdcd-9d6a-4e78-ab48-3cbb09cf108a/lib/python3.9/site-packages/flaml/ml.py in sklearn_metric_loss_score(metric_name, y_predict, y_true, labels, sample_weight, groups)
259 elif metric_name == "rmse":
260 score = np.sqrt(
--> 261 mean_squared_error(y_true, y_predict, sample_weight=sample_weight)
262 )
263 elif metric_name == "mae":

/databricks/python/lib/python3.9/site-packages/sklearn/utils/validation.py in inner_f(*args, **kwargs)
61 extra_args = len(args) - len(all_args)
62 if extra_args <= 0:
---> 63 return f(*args, **kwargs)
64
65 # extra_args > 0

/databricks/python/lib/python3.9/site-packages/sklearn/metrics/_regression.py in mean_squared_error(y_true, y_pred, sample_weight, multioutput, squared)
333 0.825...
334 """
--> 335 y_type, y_true, y_pred, multioutput = _check_reg_targets(
336 y_true, y_pred, multioutput)
337 check_consistent_length(y_true, y_pred, sample_weight)

/databricks/python/lib/python3.9/site-packages/sklearn/metrics/_regression.py in _check_reg_targets(y_true, y_pred, multioutput, dtype)
86 the dtype argument passed to check_array.
87 """
---> 88 check_consistent_length(y_true, y_pred)
89 y_true = check_array(y_true, ensure_2d=False, dtype=dtype)
90 y_pred = check_array(y_pred, ensure_2d=False, dtype=dtype)

/databricks/python/lib/python3.9/site-packages/sklearn/utils/validation.py in check_consistent_length(*arrays)
317 uniques = np.unique(lengths)
318 if len(uniques) > 1:
--> 319 raise ValueError("Found input variables with inconsistent numbers of"
320 " samples: %r" % [int(l) for l in lengths])
321

ValueError: Found input variables with inconsistent numbers of samples: [7252, 812]`