SNet2 (jax) - UnboundLocalError: local variable 'params_best' referenced before assignment

[CoteDave]

Hi,

catenets == 0.2.1

I Get this error when trying the SNet2 or the DragonNet (jax).

TNet was working #1 with the same data and parameters.

[CoteDave]

On way to avoid onp.isnan(l_curr) == True is to play with the hyperparameters of the SNet(s) (For exemple, it worked when I changed the batch_size to 10000 as the dataset contains ~500000 rows). Seems like a vanishing gradient problem....