Our custom repository of libraries had a package for pyspark which was clashing with the pyspark that is provided by the spark cluster and somehow having both works on Spark shell but does not work on a notebook.
So, renaming the pyspark library in the custom repository resolved the issue!