fix: add custom loop and docs

Signed-off-by: Mehant Kammakomati <mehant.kammakomati2@ibm.com>

Author: kmehant

plugins/online-data-mixing/artifacts/custom_loop_usage.py

@@ -91,8 +91,8 @@ class State:
     accelerator.backward(loss)
     optimizer.step()
     optimizer.zero_grad()
-    if step % 1 == 0 and accelerator.is_main_process:
-        print(f"Step {step} | Loss: {loss.item():.4f}")
+    if step_idx % 1 == 0 and accelerator.is_main_process:
+        print(f"Step {step_idx} ||| Loss: {loss.item():.4f}")
         state.log_history.append(
             {"loss": loss.item() if not torch.isnan(loss) else 1e100}
         )