@nateraw @NimaBoscarino Hey thank you for the help, I was able to mask the dataset and construct a datasetdict from it. But now there is another small issue.
I am getting this error, and not sure why this is coming. I am trying with the base code only.