fix: code refactor

kmehant · kmehant · commit 87fe0e806b0a · 2025-09-25T22:22:38.000+05:30
Signed-off-by: Mehant Kammakomati &lt;mehant.kammakomati2@ibm.com&gt;
diff --git a/plugins/online-data-mixing/src/fms_acceleration_odm/odm/dataset.py b/plugins/online-data-mixing/src/fms_acceleration_odm/odm/dataset.py
@@ -373,7 +373,7 @@ def update_sampling_weights(self, model, accelerator, state):
                     ),
                 )
                 rewards[c] += rc
-                count[c] += batch["input_ids"].shape[0]
+                count[c] += 1
             else:
                 for batch in tqdm(
                     eval_dataset_dict[self.id2cat[c]],

Original file line number	Diff line number	Diff line change
`@@ -373,7 +373,7 @@ def update_sampling_weights(self, model, accelerator, state):`
`373`	`373`	`),`
`374`	`374`	`)`
`375`	`375`	`rewards[c] += rc`
`376`		`- count[c] += batch["input_ids"].shape[0]`
	`376`	`+ count[c] += 1`
`377`	`377`	`else:`
`378`	`378`	`for batch in tqdm(`
`379`	`379`	`eval_dataset_dict[self.id2cat[c]],`