[BUG] Cannot correctly load a model trained on cpu

[KveinXu]

Cannot correctly load a model trained on cpu (using accelerator='cpu' when training).

Here is the inference code:
from deeptab import InferenceModel

infer = InferenceModel.from_path(os.path.join(LOG_DIR, "regression_model.deeptab"))
print(infer)

def predict_price(payload: dict) -> float:
X = pd.DataFrame([payload])
X_clean = infer.validate_input(X, allow_extra_columns=True)
log_pred = infer.predict(X_clean)
return float(np.exp(log_pred[0])) # invert the log transform used in training

print(predict_price(X_test.iloc[0].to_dict()))

The error log is:
Train: 3500 | Val: 750 | Test: 750
:24: ConfigWarning: TrainerConfig: lr_patience=10 >= max_epochs=5. The learning rate scheduler will never reduce the LR before training ends. Consider reducing lr_patience or increasing max_epochs.
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
💡 Tip: For seamless cloud logging and experiment tracking, try installing litlogger to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.
InferenceModel(task='regression', estimator='FTTransformerRegressor', n_features=15, features=['Unnamed: 0', 'num_0', 'num_1', ...])
💡 Tip: For seamless cloud uploads and versioning, try installing litmodels to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.
You are using a CUDA device ('NVIDIA A100-PCIE-40GB') that has Tensor Cores. To properly utilize them, you should set torch.set_float32_matmul_precision('medium' | 'high') which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision
Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/3
Train: 3500 | Val: 750 | Test: 750
:24: ConfigWarning: TrainerConfig: lr_patience=10 >= max_epochs=5. The learning rate scheduler will never reduce the LR before training ends. Consider reducing lr_patience or increasing max_epochs.
InferenceModel(task='regression', estimator='FTTransformerRegressor', n_features=15, features=['Unnamed: 0', 'num_0', 'num_1', ...])
Train: 3500 | Val: 750 | Test: 750
:24: ConfigWarning: TrainerConfig: lr_patience=10 >= max_epochs=5. The learning rate scheduler will never reduce the LR before training ends. Consider reducing lr_patience or increasing max_epochs.
InferenceModel(task='regression', estimator='FTTransformerRegressor', n_features=15, features=['Unnamed: 0', 'num_0', 'num_1', ...])
Initializing distributed: GLOBAL_RANK: 2, MEMBER: 3/3
Initializing distributed: GLOBAL_RANK: 1, MEMBER: 2/3
Segmentation fault (core dumped)