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

@@ -94,7 +94,7 @@ class State:
     loss = accelerator.gather(loss).mean()
     if step_idx % 1 == 0:
         if torch.isnan(loss):
-            raise ValueError("loss is nan")
+            loss = torch.tensor([10]) # nan -> very high loss
         print(f"Step {step_idx} ||| Loss: {loss.item():.4f}")
         state.log_history.append(
             {"loss": loss.item() if not torch.isnan(loss) else 1e100}