fix to ogbg CI test

priyakasimbeg · priyakasimbeg · commit 43d2191f4461 · 2025-08-19T16:56:15.000Z
diff --git a/tests/reference_algorithm_tests.py b/tests/reference_algorithm_tests.py
@@ -307,7 +307,7 @@ def _fake_iter():
 
         if framework == 'pytorch':
           fake_batch_iter = ogbg_input_pipeline._get_batch_iterator(
-            _fake_iter(), global_batch_size, shard=True
+            _fake_iter(), global_batch_size, shard_dim=True
           )
         else:
           fake_batch_iter = ogbg_input_pipeline._get_batch_iterator(

Original file line number	Diff line number	Diff line change
`@@ -307,7 +307,7 @@ def _fake_iter():`
`307`	`307`
`308`	`308`	`if framework == 'pytorch':`
`309`	`309`	`fake_batch_iter = ogbg_input_pipeline._get_batch_iterator(`
`310`		`- _fake_iter(), global_batch_size, shard=True`
	`310`	`+ _fake_iter(), global_batch_size, shard_dim=True`
`311`	`311`	`)`
`312`	`312`	`else:`
`313`	`313`	`fake_batch_iter = ogbg_input_pipeline._get_batch_iterator(`