Wandb

## Expected Behavior
Clean execution

## Actual Behavior
Adding the metrics logger throws wandb errors.

wandb: WARNING Tried to log to step 0 that is less than the current step 3. Steps must be monotonically increasing, so this data will be ignored. See https://wandb.me/define-metric to log data out of order.

Later has some some unexpected behavior

```
File "/home/shivajid/workshop/tensor_crunch/sample_trainer.py", line 223, in
out_data = sampler(
^^^^^^^^
File "/home/shivajid/tunixenv/lib/python3.11/site-packages/tunix/generate/sampler.py", line 705, in __call__
sampling_state = self._compiled_prefill_fn(
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/shivajid/tunixenv/lib/python3.11/site-packages/jax/_src/monitoring.py", line 97, in record_scalar
callback(event, value, **kwargs)
File "/home/shivajid/tunixenv/lib/python3.11/site-packages/tunix/sft/metrics_logger.py", line 90, in log_to_wandb
wandb.log({event: scalar_value}, step=current_step)
File "/home/shivajid/tunixenv/lib/python3.11/site-packages/wandb/sdk/lib/preinit.py", line 36, in preinit_wrapper
raise wandb.Error(f"You must call wandb.init() before {name}()")
wandb.errors.errors.Error: You must call wandb.init() before wandb.log()
```

## Steps to Reproduce the Problem

1. Add wandb to your import and login
1.
1.

## Specifications

- Version:
- Platform: TPUv6e-8