Consider adding a version that uses Tensorflow minimiser as described here https://www.tensorflow.org/probability/api_docs/python/tfp/optimizer/nelder_mead_minimize this hopefully will allow to run the minimizer on GPUs as well.