Saving tensorflow model in Databricks

is there any way how to save tensorflow model in Databricks. Im using classical aproach like in Spyder: model.save('/dbfs/FileStore/tables/mymodel.h5') Unfortunately Im still getting OOM (Out of memory error). Exactly The error is: ConnectException error: This is often caused by an OOM error that causes the connection to the Python REPL to be closed. Check your query's memory usage. Thank you for your comments. BR Tomas