[Fixed bug, but not in release] UnboundLocalError: local variable 'best_state' referenced before assignment

[Qazalbash]

I got this error while training. It is fixed, but not in the release.

Global Tuning:  79%|███████▉  | 79/100 [08:14<02:11,  6.26s/it] 
---------------------------------------------------------------------------
UnboundLocalError                         Traceback (most recent call last)
Cell In[75], line 19
      1 nf_sampler = Sampler(
      2     n_dim,
      3     rng_key_set,
   (...)
     16     use_global=True,
     17 )
---> 19 nf_sampler.sample(
     20     initial_position,
     21     data={ "data": data_set }
     22 )

File /usr/local/lib/python3.10/site-packages/flowMC/Sampler.py:196, in Sampler.sample(self, initial_position, data)
    188 last_step = initial_position
    189 for strategy in self.strategies:
    190     (
    191         rng_key,
    192         last_step,
    193         self.local_sampler,
    194         self.global_sampler,
    195         summary,
--> 196     ) = strategy(
    197         rng_key, self.local_sampler, self.global_sampler, last_step, data
    198     )
    199     self.summary[strategy.__name__] = summary

File /usr/local/lib/python3.10/site-packages/flowMC/strategy/global_tuning.py:152, in GlobalTuning.__call__(self, rng_key, local_sampler, global_sampler, initial_position, data)
    140 global_sampler.model = eqx.tree_at(
    141     lambda m: m._data_mean, global_sampler.model, data_mean
    142 )
    143 global_sampler.model = eqx.tree_at(
    144     lambda m: m._data_cov, global_sampler.model, data_cov
    145 )
    147 (
    148     rng_keys_nf,
    149     model,
    150     optim_state,
    151     loss_values,
--> 152 ) = global_sampler.model.train(
    153     rng_keys_nf,
    154     flat_chain,
    155     self.optim,
    156     self.optim_state,
    157     self.n_epochs,
    158     self.batch_size,
    159     self.verbose,
    160 )
    161 global_sampler.model = model
    162 self.optim_state = optim_state

    [... skipping hidden 1 frame]

File /usr/local/lib/python3.10/site-packages/flowMC/nfmodel/base.py:191, in NFModel.train(self, rng, data, optim, state, num_epochs, batch_size, verbose)
    188             if epoch == num_epochs:
    189                 pbar.set_description(f"Training NF, current loss: {value:.3f}")
--> 191 return rng, best_model, best_state, loss_values

UnboundLocalError: local variable 'best_state' referenced before assignment

[Qazalbash]

@kazewong It would be great if you do a quick release. Without release I will get error again on cluster.

[kazewong]

This should be fixed in the newest version now.