@Shaomu Tan​ , can you check sparktorch?

The parallel processing on Databricks clusters is mainly based on Apache Spark™. So to use the parallel processing, the library in question (PyTorch) has to be written for Spark. spark torch is an attempt to do just that.

You can also run Apache Ray on Databricks or Dask (I thought that was possible too), so bypassing Apache spark