@naman1994​  check which Databricks Runtime version you're using. I know that this works with 3.5, but didn't test with 4.0 yet.