Make bnn_mnist_sgld_blackjax.ipynb to be able to run in colab

[karm-patel]

• currently, this notebook won't run in colab, due to requirement of high memory.
• Mostly prediction loops taking much memory, we can make it batch prediction to consume less memory at a time