Patch/v2.0.2

I'm not exactly sure how it happened, but initially, I tested it with Flax version 2.0.18 on my local machine, and it was working perfectly fine. Later, I noticed that there was another version available, Flax version 2.0.19, which had been released a day after the 2.0.0 release. Consequently, I ran tests with that version as well, and it also worked flawlessly on my local setup. This is why I went ahead and created a merge request (MR) for version 2.0.1.

However, when we checked it today, it was not functioning as expected on PBR (Place-Based Research), even though it was still working in my local environment. It was only after I deleted the image and attempted to pull it again that I managed to replicate the issue.

As a result, it turns out that the correct version for use is Flax version 2.0.18, and not Flax version 2.0.19.

I have already confirmed this by testing it on https://place-based-research.cloud68.co/pbr/.