find translations using model needs some fixing

[jacksonloper]

• fix sometimes freak out if not passed tensors but numpy arrays instead
• return losses
• needs to be possible to reduce learning rate if necessary
• need to be able to supply an initial condition

[MNM-TB]

Hi! Thank you for the work with this excellent tool. I'm trying to implement it in my research but I'm running into issues completing the "basics.ipynb" examples that I think may be related to this issue.

The cell

corrections=bardensr.registration.find_translations_using_model(Xnorm,codeflat,use_tqdm_notebook=True)
runs correctly but I suspect that it isn't outputting the results correctly as I then get an error in the next cell:

plt.scatter(-bench.translations.reshape((28,3))[:,i],corrections[:,i])

The error I get is:
<class 'TypeError'> 'FindTranslationsUsingModelResults' object is not subscriptable <traceback object at 0x7e2a65ba8840>

I have had issues generating a functional conda/mamba environment but this process generates one that works for all other parts of the "basics.ipynb" notebook:

mamba create -n BarDensr_3.8 python=3.8 git ipykernel ipywidgets pandas=1.1.5 numpy=1.19.2

mamba install -c conda-forge cudatoolkit=11.0 cudnn=8.0 numpy=1.19.2 pandas=1.1.5 jaxlib==0.1.69 jax==0.2.18

export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:$LD_LIBRARY_PATH"

pip install --upgrade --no-deps git+https://github.com/jacksonloper/bardensr.git

pip install numpy==1.19.2 pandas==1.1.5 jax==0.2.18 jaxlib==0.1.69 scikit-learn==0.24 scikit-image dataclasses jaxopt==0.01 tensorflow-gpu==2.4.0 scanpy==1.4.6 seaborn==0.8.1 palantir==0.2.6  jinja2==3.1.4 protobuf==3.20 future requests-cache==1.2.1

Any suggestion on changes I can do in the environment setup to solve this would be highly appreciated.