Improve inference robustness and embedding handling

[swagat-mishra28]

Hi,

While testing the inference script locally, I noticed that the feature dimension used to initialize the model is hardcoded in inference.py.

During training, the node feature dimension is inferred directly from the dataset. Because of this, using a fixed value in the inference script could potentially lead to mismatches if the dataset or graph construction changes.

This PR updates the inference pipeline to infer the feature dimension dynamically from the dataset before initializing the model.

This keeps the inference model consistent with the training configuration and makes the script more robust when working with different datasets.

Please let me know if this approach aligns with the intended design.