Misleading print make believe Jax is loaded

[dandarm]

In the init file:

try:
print("Loaded Jax TimesFM.")
from timesfm.timesfm_jax import TimesFmJax as TimesFm
from timesfm import data_loader
except Exception as _:
print("Loaded PyTorch TimesFM.")
from timesfm.timesfm_torch import TimesFmTorch as TimesFm

can you change print rows so not to make users believe that Jax is correctly loaded, when that's not the case?
One understands that both Jax and Torch versions are loaded, while clearly is not, so now I understand hours of debugging not being able to load jax checkpoint as pointed in the finetuning notebook :)

Thanks, cheers

[rajatsen91]

Good point, sorry about that. We can change the order of the print statement, or if you want you can also submit a PR.