Update app.py

app.py

@@ -6,10 +6,21 @@ import json
 import pickle
 from pytorch_tabular import TabularModel
 from sklearn.preprocessing import LabelEncoder
+from omegaconf import OmegaConf, DictConfig
+
 
 # Load model
 model = TabularModel.load_model("FTTransformerModel")
 
+# Merge default dataloader_kwargs into config
+if isinstance(model.config, DictConfig):
+    model.config = OmegaConf.merge(
+        model.config,
+        {"dataloader_kwargs": {"batch_size": 64, "num_workers": 0}}
+    )
+else:
+    model.config.dataloader_kwargs = {"batch_size": 64, "num_workers": 0}
+
 #Encoder
 dm = model.datamodule
 if hasattr(dm, "label_encoder") and isinstance(dm.label_encoder, LabelEncoder):